Keep grants per nonce so queued and overlapping runs keep their grant

Settled grants now live for an hour from settlement instead of expiring
180 s after the request, so a run that waited in the GPU queue still
uses the grant the API already counted. Entries are keyed by nonce and
settled by the nonce in the Space-signed receipt, so a second run in the
same session no longer cancels the first. A granted entry that lapses
before use is denied with a Space-signed refund, and every failed run in
a session gets its own refund.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

access_limits.py

@@ -17,6 +17,7 @@ from quota_fallback import (
     REFUND,
     embed_parent_from,
     embed_parents,
+    grant_request_nonce,
     is_eligible_embed,
     issue_ticket,
     request_header,
@@ -30,7 +31,9 @@ logger = logging.getLogger(__name__)
 
 FREE_LIMIT_REASON = "LABS_3DAI_FREE_LIMIT"
 GRANT_TTL_SECONDS = 180
-MAX_GRANT_SESSIONS = 512
+SETTLED_GRANT_TTL_SECONDS = 3600
+MAX_GRANT_ENTRIES = 512
+MAX_GRANTS_PER_SESSION = 8
 SIGNUP_URL = "https://app.skyebrowse.com/register"
 EMBED_LIMIT_MESSAGE = (
     "You have used your free 3D AI generations. Upgrade to a Commercial plan to keep generating."
@@ -79,7 +82,9 @@ def visitor_ip(request: Any) -> str:
 @dataclass
 class GrantEntry:
     nonce: str
+    session_hash: str
     issued_at: float
+    settled_at: Optional[float] = None
     mode: Optional[str] = None
     notice: Optional[str] = None
 
@@ -93,21 +98,32 @@ class Access:
 
 
 class GrantBook:
-    def __init__(self, ttl: float = GRANT_TTL_SECONDS, max_sessions: int = MAX_GRANT_SESSIONS):
+    def __init__(
+        self,
+        ttl: float = GRANT_TTL_SECONDS,
+        max_entries: int = MAX_GRANT_ENTRIES,
+        settled_ttl: float = SETTLED_GRANT_TTL_SECONDS,
+        max_per_session: int = MAX_GRANTS_PER_SESSION,
+    ):
         self._lock = threading.Lock()
         self._entries: Dict[str, GrantEntry] = {}
+        self._lapsed: Dict[str, List[str]] = {}
         self._ttl = ttl
-        self._max_sessions = max_sessions
+        self._settled_ttl = settled_ttl
+        self._max_entries = max_entries
+        self._max_per_session = max_per_session
 
     def request(self, session_hash: str) -> Optional[Dict[str, str]]:
         if not session_hash:
             return None
         with self._lock:
             self._purge()
-            if len(self._entries) >= self._max_sessions and session_hash not in self._entries:
+            if len(self._entries) >= self._max_entries:
                 return None
-            entry = GrantEntry(uuid4().hex, time.time())
-            self._entries[session_hash] = entry
+            if len(self._session_entries(session_hash)) >= self._max_per_session:
+                return None
+            entry = GrantEntry(uuid4().hex, session_hash, time.time())
+            self._entries[entry.nonce] = entry
         ticket = issue_ticket(entry.nonce, GRANT)
         return {"ticket": ticket} if ticket else None
 
@@ -116,10 +132,13 @@ class GrantBook:
             return
         if ack.get("ticket") != receipt.get("ticket"):
             return
+        nonce = grant_request_nonce(receipt.get("ticket"))
         with self._lock:
-            entry = self._entries.get(session_hash)
-            if entry is None or entry.mode or entry.notice:
+            self._purge()
+            entry = self._entries.get(nonce) if nonce else None
+            if entry is None or entry.session_hash != session_hash or entry.settled_at is not None:
                 return
+            entry.settled_at = time.time()
             grant = verify_grant(ack.get("grant"), entry.nonce)
             if grant is not None:
                 entry.mode = grant["mode"]
@@ -131,18 +150,45 @@ class GrantBook:
     def take(self, session_hash: str) -> Access:
         with self._lock:
             self._purge()
-            entry = self._entries.pop(session_hash, None)
-        if entry is None or (entry.mode is None and entry.notice is None):
-            return Access(mode="denied", notice=embed_notice(EMBED_UNAVAILABLE_MESSAGE))
-        if entry.mode is None:
-            return Access(mode="denied", notice=embed_notice(entry.notice))
+            entry = self._pop_settled(session_hash)
+            lapsed = None if entry else self._pop_lapsed(session_hash)
+        if lapsed:
+            return Access(mode="denied", nonce=lapsed, notice=embed_notice(EMBED_UNAVAILABLE_MESSAGE))
+        if entry is None or entry.mode is None:
+            notice = entry.notice if entry and entry.notice else EMBED_UNAVAILABLE_MESSAGE
+            return Access(mode="denied", notice=embed_notice(notice))
         return Access(nonce=entry.nonce)
 
+    def _pop_settled(self, session_hash: str) -> Optional[GrantEntry]:
+        settled = [entry for entry in self._session_entries(session_hash) if entry.settled_at is not None]
+        return self._entries.pop(min(settled, key=lambda entry: entry.settled_at).nonce) if settled else None
+
+    def _pop_lapsed(self, session_hash: str) -> Optional[str]:
+        lapsed = self._lapsed.get(session_hash)
+        if not lapsed:
+            return None
+        nonce = lapsed.pop(0)
+        if not lapsed:
+            del self._lapsed[session_hash]
+        return nonce
+
+    def _session_entries(self, session_hash: str) -> List[GrantEntry]:
+        return [entry for entry in self._entries.values() if entry.session_hash == session_hash]
+
     def _purge(self) -> None:
-        cutoff = time.time() - self._ttl
-        for key, entry in list(self._entries.items()):
-            if entry.issued_at < cutoff:
-                del self._entries[key]
+        now = time.time()
+        for nonce, entry in list(self._entries.items()):
+            pending = entry.settled_at is None
+            started = entry.issued_at if pending else entry.settled_at
+            if started < now - (self._ttl if pending else self._settled_ttl):
+                del self._entries[nonce]
+                if entry.mode is not None:
+                    self._lapse(entry)
+
+    def _lapse(self, entry: GrantEntry) -> None:
+        self._lapsed.setdefault(entry.session_hash, []).append(entry.nonce)
+        while len(self._lapsed) > self._max_entries:
+            del self._lapsed[next(iter(self._lapsed))]
 
 
 class PublicAllowance:
@@ -256,12 +302,15 @@ class AccessPolicy:
         self.grants = grants
         self.public = public
         self._notices: Dict[str, str] = {}
-        self._refunds: Dict[str, str] = {}
+        self._refunds: Dict[str, List[str]] = {}
         self._lock = threading.Lock()
 
     def authorize(self, request: Any, session_hash: str) -> Access:
         if grants_required(request):
-            return self.grants.take(session_hash)
+            access = self.grants.take(session_hash)
+            if access.mode == "denied":
+                self.generation_failed(session_hash, access)
+            return access
         if not (self.public.enabled and is_browser_visit(request)):
             return Access()
         key = self.public.key_for(request)
@@ -278,11 +327,14 @@ class AccessPolicy:
         self.public.refund(access.visitor_key)
         if access.nonce:
             with self._lock:
-                self._refunds[session_hash] = access.nonce
+                self._refunds.setdefault(session_hash, []).append(access.nonce)
 
     def refund_receipt(self, session_hash: str) -> Optional[Dict[str, str]]:
         with self._lock:
-            nonce = self._refunds.pop(session_hash, None)
+            pending = self._refunds.get(session_hash) or []
+            nonce = pending.pop(0) if pending else None
+            if not pending:
+                self._refunds.pop(session_hash, None)
         ticket = issue_ticket(nonce, REFUND) if nonce else None
         return {"ticket": ticket} if ticket else None
 

quota_fallback.py

@@ -208,6 +208,14 @@ def verify_permit(token: Any, receipt: Dict[str, Any]) -> Optional[Dict[str, Any
     return payload
 
 
+def grant_request_nonce(token: Any) -> Optional[str]:
+    payload = _open_document(token)
+    if payload is None or payload.get("operation") != GRANT or payload.get("space") != space_identifier():
+        return None
+    nonce = payload.get("nonce")
+    return nonce if isinstance(nonce, str) else None
+
+
 def verify_grant(token: Any, nonce: str) -> Optional[Dict[str, Any]]:
     payload = _open_document(token)
     if payload is None or payload.get("purpose") != GRANT_PURPOSE:

tests/test_access_limits.py

@@ -107,6 +107,57 @@ class GrantBookTests(unittest.TestCase):
             self.book.settle("session-a", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
             self.assertEqual(self.book.take("session-a").mode, "denied")
 
+    def test_grant_requests_are_bounded_per_session_and_in_total(self):
+        for _ in range(al.MAX_GRANTS_PER_SESSION):
+            self.assertIsNotNone(self.book.request("session-a"))
+        self.assertIsNone(self.book.request("session-a"))
+        self.assertIsNotNone(self.book.request("session-b"))
+        book = al.GrantBook(max_entries=1)
+        self.assertIsNotNone(book.request("one"))
+        self.assertIsNone(book.request("two"))
+
+    def settle(self, receipt, nonce, session="session-a"):
+        self.book.settle(session, {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
+
+    def test_a_settled_grant_survives_a_long_gpu_queue(self):
+        receipt, nonce = self.issue()
+        self.settle(receipt, nonce)
+        with mock.patch.object(al.time, "time", return_value=time.time() + al.GRANT_TTL_SECONDS + 60):
+            access = self.book.take("session-a")
+        self.assertEqual((access.mode, access.nonce), ("native", nonce))
+
+    def test_a_lapsed_grant_is_denied_with_its_nonce_for_a_refund(self):
+        receipt, nonce = self.issue()
+        self.settle(receipt, nonce)
+        with mock.patch.object(al.time, "time", return_value=time.time() + al.SETTLED_GRANT_TTL_SECONDS + 1):
+            access = self.book.take("session-a")
+        self.assertEqual((access.mode, access.nonce), ("denied", nonce))
+        self.assertIsNone(self.book.take("session-a").nonce)
+
+    def test_overlapping_runs_in_one_session_each_keep_their_grant(self):
+        first, first_nonce = self.issue()
+        second, second_nonce = self.issue()
+        self.settle(first, first_nonce)
+        self.settle(second, second_nonce)
+        self.assertEqual(self.book.take("session-a").nonce, first_nonce)
+        self.assertEqual(self.book.take("session-a").nonce, second_nonce)
+        self.assertEqual(self.book.take("session-a").mode, "denied")
+
+    def test_overlapping_runs_keep_their_own_refusals(self):
+        first, first_nonce = self.issue()
+        second, _nonce = self.issue()
+        self.settle(first, first_nonce)
+        self.book.settle("session-a", {"ticket": second["ticket"], "reason": al.FREE_LIMIT_REASON, "error": ""}, second)
+        self.assertEqual(self.book.take("session-a").nonce, first_nonce)
+        self.assertIn(al.EMBED_LIMIT_MESSAGE, self.book.take("session-a").notice)
+
+    def test_an_ack_for_an_unknown_nonce_marks_no_other_grant(self):
+        receipt, nonce = self.issue()
+        stranger = {"ticket": qf.issue_ticket("e" * 32, qf.GRANT)}
+        self.book.settle("session-a", {"ticket": stranger["ticket"], "grant": grant_for("e" * 32)}, stranger)
+        self.settle(receipt, nonce)
+        self.assertEqual(self.book.take("session-a").nonce, nonce)
+
 
 class PublicAllowanceTests(unittest.TestCase):
     def setUp(self):
@@ -211,6 +262,29 @@ class AccessPolicyTests(unittest.TestCase):
         self.assertIsNone(qf.verify_receipt(refund["ticket"]))
         self.assertIsNone(self.policy.refund_receipt("e"))
 
+    def grant(self, session="e"):
+        receipt = self.policy.grants.request(session)
+        nonce = opened(receipt["ticket"])["nonce"]
+        self.policy.grants.settle(session, {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
+        return nonce
+
+    def test_overlapping_failed_generations_are_each_refunded(self):
+        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}):
+            nonces = [self.grant(), self.grant()]
+            for _ in nonces:
+                access = self.policy.authorize(embed_request(), "e")
+                self.assertEqual(access.mode, "native")
+                self.policy.generation_failed("e", access)
+        self.assertEqual([opened(self.policy.refund_receipt("e")["ticket"])["nonce"] for _ in nonces], nonces)
+        self.assertIsNone(self.policy.refund_receipt("e"))
+
+    def test_a_lapsed_grant_is_refunded_when_its_run_is_denied(self):
+        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}):
+            nonce = self.grant()
+            with mock.patch.object(al.time, "time", return_value=time.time() + al.SETTLED_GRANT_TTL_SECONDS + 1):
+                self.assertEqual(self.policy.authorize(embed_request(), "e").mode, "denied")
+        self.assertEqual(opened(self.policy.refund_receipt("e")["ticket"])["nonce"], nonce)
+
     def test_notices_are_delivered_once(self):
         self.policy.remember("s", "<div>notice</div>")
         self.assertEqual(self.policy.pop_notice("s"), "<div>notice</div>")