fix: Make ticket claim atomic.

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(

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
         )