diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index 22704334..33ed853d 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -88,15 +88,19 @@ def check_jwt_ro(cfg: BeamlineConfig, data: TokenData) -> None: def check_jwt_rw(cfg: BeamlineConfig, data: TokenData) -> None: check_jwt_ro(cfg, data) + holder = cfg.baton_holder + if holder is None or holder.session != data.session: + raise UserRightsException( + message="You do not hold the baton.", + status_code=403, + code=AuthErrorCode.NOT_BATON_HOLDER, + ) + try: - cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS) + cfg.try_extend_active_session(data.session, SESSION_EXPIRE_SECONDS) except Exception as e: - raise AuthenticationException( - message="Another session is active.", - status_code=401, - headers={"WWW-Authenticate": "Bearer"}, - code=AuthErrorCode.SESSION_ALREADY_ACTIVE, - ) from e + # In case something is wrong but you are holder (maybe redis expiry?) + cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS) def check_jwt_staff_only(data: TokenData) -> None: if not data.staff: @@ -108,15 +112,18 @@ def check_jwt_staff_only(data: TokenData) -> None: def check_jwt_staff(cfg: BeamlineConfig, data: TokenData) -> None: check_jwt_staff_only(data) + + holder = cfg.baton_holder + if holder and holder.session != data.session: + # Staff can take over if they don't have it, but they need to use force_current_session + # or request_baton (which does staff override). + # If they are calling a RW endpoint, they SHOULD already be the holder. + pass + try: + cfg.try_extend_active_session(data.session, SESSION_EXPIRE_SECONDS) + except Exception: cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS) - except Exception as e: - raise AuthenticationException( - message="Another session is active.", - status_code=401, - headers={"WWW-Authenticate": "Bearer"}, - code=AuthErrorCode.SESSION_ALREADY_ACTIVE, - ) from e def force_current_sesion(cfg: BeamlineConfig, data: TokenData) -> None: cfg.execute_baton_transfer( @@ -222,6 +229,7 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: session_state = cfg.session_state(data.session) if session_state == SessionsStateEnum.Vacant: + # State: VACANT -> OWNED (B) cfg.execute_baton_transfer( to_session=data.session, to_username=data.sub, @@ -232,19 +240,27 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: return {"granted": True, "message": "Baton acquired (beamline was vacant)"} if session_state == SessionsStateEnum.OwnedByYou: - cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS) + cfg.try_extend_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 session_state == SessionsStateEnum.PendingYouToElse: + # Already requested, waiting for holder response or timeout + existing_request = cfg.pending_baton_request + if existing_request: + elapsed = time.time() - existing_request.created_at + 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)", + } + holder = cfg.baton_holder + # Handle Staff override if data.staff: if not cfg.can_transfer_baton_now(): + # State: OWNED (A) -> PENDING (A -> B, B=staff) -> Queue -> OWNED (B) cfg.queued_baton_transfer = BatonTransferQueue( target_session=data.session, target_username=data.sub, @@ -258,6 +274,7 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: "message": "Staff override queued - will transfer when beamline is available", } + # State: OWNED (A) -> OWNED (B, B=staff) cfg.execute_baton_transfer( to_session=data.session, to_username=data.sub, @@ -267,24 +284,25 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: ) 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)", - } + # Policy check: non-staff requesting 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": "Another user already has a pending request", + "message": "Requesting baton from staff is disabled by backend policy.", } + # State: OWNED (A) -> PENDING (A -> B) + existing_request = cfg.pending_baton_request + if existing_request and existing_request.status == BatonRequestStatus.PENDING: + # Someone else already has a request pending + return { + "error": True, + "message": f"Another user ({existing_request.requester_username}) already has a pending request", + } + + # Beamline busy notification + is_busy = not cfg.can_transfer_baton_now() + request = BatonRequest( request_id=str(uuid.uuid4()), requester_username=data.sub, @@ -298,11 +316,16 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: ) cfg.set_pending_baton_request(request, timeout_sec=BATON_REQUEST_TIMEOUT_SECONDS) + msg = f"Request sent to {holder.username if holder else 'current holder'}" + if is_busy: + msg += " (Note: beamline is currently busy, transfer will be queued if accepted)" + 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'}", + "message": msg, + "beamline_busy": is_busy, } def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) -> dict: @@ -320,6 +343,7 @@ def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) return {"error": True, "message": "No pending request to respond to"} if accept: + # State: PENDING (A -> B) -> OWNED (B) if cfg.can_transfer_baton_now(): cfg.execute_baton_transfer( to_session=pending.requester_session, @@ -330,6 +354,7 @@ def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) ) return {"accepted": True, "transferred": True, "message": "Baton transferred"} else: + # Beamline busy, queue the transfer cfg.queued_baton_transfer = BatonTransferQueue( target_session=pending.requester_session, target_username=pending.requester_username, @@ -342,12 +367,13 @@ def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) return { "accepted": True, "queued": True, - "message": "Request accepted - will transfer when beamline is available" + "message": "Request accepted. Baton will be transferred as soon as beamline is available.", } else: + # State: PENDING (A -> B) -> OWNED (A) pending.status = BatonRequestStatus.REFUSED - cfg.set_pending_baton_request(pending, timeout_sec=5) - return {"refused": True, "message": "Request refused"} + cfg.set_pending_baton_request(pending, timeout_sec=10) + return {"refused": True, "message": "Baton request refused"} def release_baton(cfg: BeamlineConfig, data: TokenData) -> dict: """ diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 3011466f..482a5605 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -115,29 +115,31 @@ class BeamlineConfig: Returns: int if present and valid, otherwise None. """ - raw = self.__client.get(f"{self.__bl}:active_session") - if raw is None: - return None - try: - return int(raw) - except (TypeError, ValueError): - logger.warning( - "Invalid active_session value in redis; treating as missing", - extra={"beamline": self.__bl, "raw": raw}, - ) + baton = self.baton_holder + if baton is None: return None + return baton.session def session_status(self, session: int) -> SessionStatus: return SessionStatus(session=self.session_state(session), current_pgroup=self.pgroup) def session_state(self, session: int) -> SessionsStateEnum: - curr_session = self.active_session - if curr_session is None: + holder = self.baton_holder + pending = self.pending_baton_request + + if holder is None: return SessionsStateEnum.Vacant - elif curr_session == session: + + if holder.session == session: + # You are the holder. Check if someone else requested from you. + if pending and pending.status == BatonRequestStatus.PENDING and pending.holder_session == session: + return SessionsStateEnum.PendingElseToYou return SessionsStateEnum.OwnedByYou else: + # Someone else is the holder. Check if you requested from them. + if pending and pending.status == BatonRequestStatus.PENDING and pending.requester_session == session: + return SessionsStateEnum.PendingYouToElse return SessionsStateEnum.OwnedByElse def try_set_active_session(self, session: int, expiry_sec: int) -> None: @@ -776,4 +778,5 @@ class BeamlineConfig: if __name__ == "__main__": from aare.common.beamline import mx_beamline cfg = BeamlineConfig(bl=mx_beamline()) - cfg.allow_non_staff_request_from_staff = True \ No newline at end of file + cfg.allow_non_staff_request_from_staff = True + cfg.state_busy = False \ No newline at end of file