From 72904b0b9598d0ea1527c4dab4d221e866da4172 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 14 Sep 2026 14:06:59 +0100 Subject: [PATCH] fix: Make ticket claim atomic. --- docsgpt/api/devices/session_events.py | 15 ++++-------- docsgpt/devices/broker.py | 35 +++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/docsgpt/api/devices/session_events.py b/docsgpt/api/devices/session_events.py index afb137e3..65b47f2f 100644 --- a/docsgpt/api/devices/session_events.py +++ b/docsgpt/api/devices/session_events.py @@ -13,7 +13,7 @@ import json import logging import time from functools import partial -from typing import AsyncIterator, Optional +from typing import AsyncIterator import anyio from starlette.requests import Request @@ -63,13 +63,6 @@ def _signed_path(request: Request) -> str: return path -def _claim_session(broker: DeviceBroker, device: dict, session_id: str) -> Optional[SessionState]: - """Open the device's session iff ``session_id`` is its unexpired ticket; runs in a worker thread.""" - if not broker.validate_ticket(device["id"], session_id): - return None - return broker.register_session(device["id"], device["user_id"]) - - async def _close_session(broker: DeviceBroker, session_id: str) -> None: """Close the session once its response is over.""" broker.close_session(session_id, reason="stream_end") @@ -132,11 +125,13 @@ async def device_session_events(request: Request) -> Response: return _error(*failure) log_context.bind(user_id=device["user_id"]) broker = get_broker() - # Claiming consumes the ticket, so it happens before the response opens and + # Redeeming consumes the ticket, so it happens before the response opens and # ``on_close`` closes the session even if the body never starts. Shielded # so a cancellation can't leave a registered session with no response. with anyio.CancelScope(shield=True): - sess = await anyio.to_thread.run_sync(_claim_session, broker, device, session_id) + sess = await anyio.to_thread.run_sync( + broker.redeem_ticket, device["id"], device["user_id"], session_id + ) if sess is None: return _error("session_ticket_invalid", 410) return ClosingStreamingResponse( diff --git a/docsgpt/devices/broker.py b/docsgpt/devices/broker.py index 781df24b..09d55360 100644 --- a/docsgpt/devices/broker.py +++ b/docsgpt/devices/broker.py @@ -61,6 +61,16 @@ def _ticket_key(device_id: str) -> str: return f"dev:ticket:{device_id}" +# Deletes the device's ticket iff it still holds the presented one. Runs +# server-side so two requests racing with one ticket can't both redeem it. +_REDEEM_TICKET_LUA = """ +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +""" + + def _inv_key(invocation_id: str) -> str: return f"dev:inv:{invocation_id}" @@ -142,6 +152,31 @@ class DeviceBroker: except Exception: logger.exception("ticket lookup failed for %s", device_id) session_id = _as_str(issued) if issued else f"st_{uuid.uuid4().hex}" + return self._adopt_session(session_id, device_id, user_id) + + def redeem_ticket(self, device_id: str, user_id: str, ticket: str) -> Optional[SessionState]: + """Consume the device's poll-issued ``ticket`` and open its session under it. + + Returns ``None``, and opens nothing, unless ``ticket`` is still the + device's unexpired ticket. The compare and delete are one Redis step, + so a replay racing the first redeem can't replace its session. + """ + if not ticket: + return None + redis = get_redis_instance() + if redis is None: + return None + try: + redeemed = redis.eval(_REDEEM_TICKET_LUA, 1, _ticket_key(device_id), ticket) + except Exception: + logger.exception("ticket redeem failed for %s", device_id) + return None + if not redeemed: + return None + return self._adopt_session(ticket, device_id, user_id) + + def _adopt_session(self, session_id: str, device_id: str, user_id: str) -> SessionState: + """Make a new session the device's live one, closing the one it replaces.""" sess = SessionState( session_id=session_id, device_id=device_id, user_id=user_id )