BATON: tried to convert from sessions to baton. This appears to be causign many issues. need to unpack them on Monday.

This commit is contained in:
2026-03-27 17:09:05 +01:00
parent 53608f917a
commit 7ae9889606
2 changed files with 83 additions and 54 deletions
+66 -40
View File
@@ -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:
"""
+17 -14
View File
@@ -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
cfg.allow_non_staff_request_from_staff = True
cfg.state_busy = False