refactor(free-tier): one per-profile challenge record behind one lock

_pending (dict-of-dicts keyed "payload"/"deadline"/"outcome"),
_work_locks and _gave_up_until were three per-profile module globals,
two of them mutated outside the lock. They become one _ProfileChallenge
dataclass in a single _profiles map behind _state_lock; the record holds
the BrowserChallenge itself, so lookups are attribute reads and
as_payload() runs only when a status read asks for it. _set_pending
becomes _record(challenge, new_attempt=...). The test-only _presenter
seam goes: tests patch present() itself. _sleep stays, same idiom as
free_tier_bootstrap.

One edge differs, for the better: an EXPIRED same-URL record is now
treated as a new ticket (fresh deadline) instead of re-recorded with
min(new, old-expired) deadline, i.e. born expired.
This commit is contained in:
kshitijk4poor
2026-10-06 16:34:29 +05:30
committed by kshitij
parent 50b0082ed0
commit fa6ef8ba6d
3 changed files with 97 additions and 91 deletions

View File

@@ -39,7 +39,7 @@ import platform
import sys
import threading
import time
from dataclasses import dataclass, replace
from dataclasses import dataclass, field, replace
from typing import Any, Callable, Dict, Optional, TypeVar
from urllib.parse import urlparse
@@ -60,6 +60,9 @@ _POLL_MIN_SECONDS, _POLL_MAX_SECONDS = 1.0, 5.0
# ``expires_in`` is network input that ends up in timers on two runtimes: keep it finite and sane.
_EXPIRES_MIN_SECONDS, _EXPIRES_MAX_SECONDS, _EXPIRES_DEFAULT_SECONDS = 30.0, 900.0, 600.0
_MESSAGE_MAX_CHARS = 300
# After one attempt ran out of patience, the callers queued behind it fail fast for this long
# instead of each parking for another full wait (boot fires a burst of token reads).
_GAVE_UP_COOLDOWN_SECONDS = 30.0
CHALLENGE_COPY = "Nous needs to run a quick check before starting your free session."
@@ -192,12 +195,23 @@ def challenge_error(payload: Dict[str, Any], *, portal_base_url: str, anon_token
auth_state=auth_state)
# --- Pending challenge: what a status read shows --------------------------------------------------------
# --- Per-profile state: what a status read shows, and who is working it ---------------------------------
_pending_lock = threading.Lock()
_pending: Dict[str, Dict[str, Any]] = {}
# One challenge worked at a time PER PROFILE: a sibling profile's wait must not queue behind it.
_work_locks: Dict[str, threading.Lock] = {}
@dataclass
class _ProfileChallenge:
"""One profile's challenge state, behind ``_state_lock``. ``work_lock`` serialises waits PER
PROFILE (a sibling profile's wait must not queue behind it)."""
work_lock: threading.Lock = field(default_factory=threading.Lock)
pending: Optional[BrowserChallenge] = None
deadline: float = 0.0
outcome: Optional[str] = None # how the host's window ended, for ``pending``'s attempt
gave_up_until: float = 0.0
_state_lock = threading.Lock()
_profiles: Dict[str, _ProfileChallenge] = {}
_opened_urls: set[str] = set() # challenge URLs whose browser tab was opened (one per ticket)
# Set by :func:`background_caller`: this caller has nobody waiting on it.
_background: contextvars.ContextVar[bool] = contextvars.ContextVar("anon_challenge_background", default=False)
@@ -214,83 +228,90 @@ def background_caller():
_background.reset(token)
def _work_lock() -> threading.Lock:
with _pending_lock:
return _work_locks.setdefault(_profile_key(), threading.Lock())
def _profile() -> _ProfileChallenge:
"""This profile's record. Callers hold ``_state_lock`` for anything but ``work_lock``."""
with _state_lock:
return _profiles.setdefault(_mint_memo_key(), _ProfileChallenge())
def _profile_key() -> str:
return _mint_memo_key()
def _live(record: _ProfileChallenge) -> Optional[BrowserChallenge]:
if record.pending and record.deadline <= time.monotonic():
record.pending, record.outcome = None, None
return record.pending
def pending_challenge() -> Optional[Dict[str, Any]]:
"""The challenge this profile is waiting on, for ``free_tier.status`` (a client that connected
after the event fired still learns it has a window to open). None once cleared or expired, or
once the host said how its window ended (a window the user closed must not come back)."""
with _pending_lock:
entry = _pending.get(_profile_key())
if entry and entry["deadline"] <= time.monotonic():
_pending.pop(_profile_key(), None)
entry = None
if not entry or entry["outcome"] not in (None, "timeout", "error"):
record = _profile()
with _state_lock:
challenge = _live(record)
if challenge is None or record.outcome not in (None, "timeout", "error"):
return None
return {**entry["payload"], "expires_in": math.ceil(entry["deadline"] - time.monotonic())}
return {**challenge.as_payload(), "expires_in": math.ceil(record.deadline - time.monotonic())}
def _clear_pending() -> None:
with _pending_lock:
_pending.pop(_profile_key(), None)
record = _profile()
with _state_lock:
record.pending, record.outcome = None, None
def _set_pending(challenge: BrowserChallenge, *, new_attempt: bool = False) -> BrowserChallenge:
"""Record *challenge* as the profile's pending one. A replay of the same ticket keeps the
original deadline and (unless *new_attempt*) its host outcome; the attempt number is what
lets a fresh foreground attempt reopen a window the user closed."""
with _pending_lock:
key = _profile_key()
previous = _pending.get(key)
if previous and previous["payload"]["url"] != challenge.url:
previous = None
attempt = previous["payload"]["attempt"] + new_attempt if previous else 0
def _record(challenge: BrowserChallenge, *, new_attempt: bool) -> BrowserChallenge:
"""Record *challenge* as the profile's pending one. A replay of the same ticket (a background
read, a second 428) keeps the original deadline, attempt and host outcome; ``new_attempt`` (a
foreground caller about to wait) bumps the attempt, which lets it reopen a window the user closed."""
record = _profile()
with _state_lock:
previous = _live(record)
deadline = time.monotonic() + challenge.expires_in
if previous:
deadline = min(deadline, previous["deadline"])
challenge = replace(challenge, attempt=attempt)
_pending[key] = {"payload": challenge.as_payload(), "deadline": deadline,
"outcome": previous["outcome"] if previous and not new_attempt else None}
return challenge
def _matching_entry(url: str, attempt: int) -> Optional[Dict[str, Any]]:
entry = _pending.get(_profile_key())
if entry and entry["payload"]["url"] == url and entry["payload"]["attempt"] == attempt:
return entry
return None
if previous is None or previous.url != challenge.url:
record.pending, record.deadline, record.outcome = replace(challenge, attempt=0), deadline, None
else:
attempt = previous.attempt + 1 if new_attempt else previous.attempt
record.pending, record.deadline = replace(challenge, attempt=attempt), min(deadline, record.deadline)
if new_attempt:
record.outcome = None
return record.pending
def record_host_outcome(url: str, attempt: int, outcome: str) -> bool:
"""A presentation hint: wakes a waiter to re-mint, never grants clearance locally. A late
reply belongs to the attempt that opened the window, not to a newer retry."""
with _pending_lock:
entry = _matching_entry(url, attempt)
if not entry or entry["deadline"] <= time.monotonic():
record = _profile()
with _state_lock:
challenge = _live(record)
if challenge is None or (challenge.url, challenge.attempt) != (url, attempt):
return False
entry["outcome"] = outcome
record.outcome = outcome
return True
def _host_outcome(challenge: BrowserChallenge) -> Optional[str]:
with _pending_lock:
entry = _matching_entry(challenge.url, challenge.attempt)
return entry["outcome"] if entry else None
record = _profile()
with _state_lock:
pending = record.pending
matches = pending is not None and (pending.url, pending.attempt) == (challenge.url, challenge.attempt)
return record.outcome if matches else None
def _give_up() -> None:
record = _profile()
with _state_lock:
record.gave_up_until = time.monotonic() + _GAVE_UP_COOLDOWN_SECONDS
def _gave_up_recently() -> bool:
record = _profile()
with _state_lock:
return time.monotonic() < record.gave_up_until
def reset_for_tests() -> None:
with _pending_lock:
_pending.clear()
_work_locks.clear()
_gave_up_until.clear()
_opened_urls.clear()
with _state_lock:
_profiles.clear()
_opened_urls.clear()
# --- Presenting ---------------------------------------------------------------------------------------
@@ -311,9 +332,6 @@ def _announce(challenge: BrowserChallenge) -> bool:
return False
_opened_urls: set = set() # challenge URLs whose browser tab was opened
def _present_in_terminal(challenge: BrowserChallenge) -> None:
from hermes_cli.auth_device_flow import _can_open_graphical_browser, _is_remote_session
opened = False
@@ -336,7 +354,7 @@ def _present_in_terminal(challenge: BrowserChallenge) -> None:
def present(challenge: BrowserChallenge) -> None:
"""Get the URL in front of something that can load it. Seam for tests (``_presenter``).
"""Get the URL in front of something that can load it.
The desktop backend hands every challenge to its hidden window. The stdio TUI gateway (whose
stderr the TUI keeps as a log, never shows) hands its client a required one a foreground caller
@@ -351,10 +369,6 @@ def present(challenge: BrowserChallenge) -> None:
_present_in_terminal(challenge)
_presenter: Callable[[BrowserChallenge], None] = present
_sleep = time.sleep # seam for tests
# --- Working a challenge ----------------------------------------------------------------------------
@@ -380,14 +394,7 @@ def _poll_status(client: httpx.Client, portal_base_url: str, anon_token: str) ->
_WAITING = ("pending", "needs_interaction")
# After one attempt ran out of patience, the callers queued behind it fail fast for this long
# instead of each parking for another full wait (boot fires a burst of token reads).
_GAVE_UP_COOLDOWN_SECONDS = 30.0
_gave_up_until: Dict[str, float] = {}
def _give_up(key: str) -> None:
_gave_up_until[key] = time.monotonic() + _GAVE_UP_COOLDOWN_SECONDS
_sleep = time.sleep # seam for tests (same idiom as ``free_tier_bootstrap._sleep``)
def wait_for_challenge(exc: AnonChallengeRequired) -> bool:
@@ -398,8 +405,7 @@ def wait_for_challenge(exc: AnonChallengeRequired) -> bool:
already settled on its first status read and returns without presenting anything."""
from hermes_cli.auth import _resolve_verify
from hermes_cli.auth_nous import _nous_http_client
key = _profile_key()
if time.monotonic() < _gave_up_until.get(key, 0.0):
if _gave_up_recently():
return False
challenge = exc.challenge
deadline = time.monotonic() + min(CHALLENGE_WAIT_SECONDS, challenge.expires_in)
@@ -410,8 +416,8 @@ def wait_for_challenge(exc: AnonChallengeRequired) -> bool:
with _nous_http_client(10.0, verify) as client:
if _poll_status(client, exc.portal_base_url, exc.anon_token) not in _WAITING:
return True
challenge = exc.challenge = _set_pending(challenge, new_attempt=True)
_presenter(challenge)
challenge = exc.challenge = _record(challenge, new_attempt=True)
present(challenge)
while time.monotonic() < deadline:
if _host_outcome(challenge) is not None:
return True
@@ -423,7 +429,7 @@ def wait_for_challenge(exc: AnonChallengeRequired) -> bool:
told_interactive = True
if client_surface() != "desktop":
print(" Finish the quick check in your browser to continue.", file=sys.stderr)
_give_up(key)
_give_up()
return False
@@ -452,9 +458,9 @@ def run_with_challenge(exchange: Callable[[], T]) -> T:
# A messaging gateway has nobody at its console to clear a check, and its token reads can run
# on the event loop: it never waits either.
if _background.get() or client_surface() == "gateway":
_presenter(_set_pending(first.challenge))
present(_record(first.challenge, new_attempt=False))
raise _still_pending(first)
with _work_lock():
with _profile().work_lock:
wait_for_challenge(first)
# Even a timed-out status poll can lag a committed clearance. Mint is
# authoritative and gets one final attempt before we report a pending check.
@@ -463,8 +469,8 @@ def run_with_challenge(exchange: Callable[[], T]) -> T:
except AnonChallengeRequired as again:
# Whatever ended the wait (status, a host hint, patience), a second 428 starts the fail-fast
# cooldown; the attempt after it opens a fresh window (``new_attempt``).
again.challenge = _set_pending(again.challenge)
_give_up(_profile_key())
again.challenge = _record(again.challenge, new_attempt=False)
_give_up()
raise _still_pending(again)
except AuthError as error:
if not error.retryable:
@@ -480,4 +486,4 @@ def note_optional_challenges(payload: Dict[str, Any], portal_base_url: str) -> N
if challenge is None:
_clear_pending()
elif not challenge.required and client_surface() == "desktop":
_announce(_set_pending(challenge))
_announce(_record(challenge, new_attempt=False))

View File

@@ -22,7 +22,7 @@ def nas(monkeypatch, tmp_path):
@pytest.fixture
def presented(monkeypatch):
seen: list[anon_challenge.BrowserChallenge] = []
monkeypatch.setattr(anon_challenge, "_presenter", seen.append)
monkeypatch.setattr(anon_challenge, "present", seen.append)
return seen
@@ -69,7 +69,7 @@ class TestRequiredChallenge:
def failed(challenge):
assert anon_challenge.record_host_outcome(challenge.url, challenge.attempt, "error")
monkeypatch.setattr(anon_challenge, "_presenter", failed)
monkeypatch.setattr(anon_challenge, "present", failed)
with pytest.raises(anon_auth.AuthError, match="couldn't finish"):
_resolve()
assert [p for _, p in nas.calls].count("/api/anonymous/token") == 2
@@ -157,14 +157,14 @@ class TestRequiredChallenge:
def present(_challenge):
with _nous_shared_store_lock(timeout_seconds=0.5):
pass
monkeypatch.setattr(anon_challenge, "_presenter", present)
monkeypatch.setattr(anon_challenge, "present", present)
assert _resolve()["api_key"]
def test_status_shows_the_pending_challenge_while_it_is_worked(self, nas, monkeypatch):
nas.challenge_required = True
nas.challenge_statuses = ["pending"]
seen = []
monkeypatch.setattr(anon_challenge, "_presenter", lambda _c: seen.append(anon_challenge.pending_challenge()))
monkeypatch.setattr(anon_challenge, "present", lambda _c: seen.append(anon_challenge.pending_challenge()))
_resolve()
assert seen == [{"type": "browser", "url": nas.challenge_url, "required": True,
"expires_in": 600, "message": "A quick check first.", "attempt": 0}]
@@ -381,11 +381,11 @@ class TestOptionalChallenge:
clock = [100.0]
monkeypatch.setattr(anon_challenge.time, "monotonic", lambda: clock[0])
challenge = anon_challenge.BrowserChallenge(nas.challenge_url, False, 60, 2, "check")
anon_challenge._set_pending(challenge)
anon_challenge._record(challenge, new_attempt=False)
clock[0] += 20
anon_challenge._set_pending(challenge)
anon_challenge._record(challenge, new_attempt=False)
assert anon_challenge.pending_challenge()["expires_in"] == 40
retry = anon_challenge._set_pending(challenge, new_attempt=True)
retry = anon_challenge._record(challenge, new_attempt=True)
assert not anon_challenge.record_host_outcome(challenge.url, 0, "done")
assert anon_challenge.pending_challenge()["attempt"] == retry.attempt
clock[0] += 41

View File

@@ -85,7 +85,7 @@ def test_status_carries_the_pending_browser_challenge_and_drops_it_once_cleared(
challenge = anon_challenge.BrowserChallenge(
"https://portal.example.test/challenge?code=t", True, 600, 2, "A quick check first.")
anon_challenge._set_pending(challenge)
anon_challenge._record(challenge, new_attempt=False)
try:
assert _call("free_tier.status")["challenge"] == challenge.as_payload()
finally:
@@ -98,10 +98,10 @@ def test_window_outcomes_are_scoped_and_do_not_grant_auth(guest, tmp_path):
from hermes_constants import set_hermes_home_override, reset_hermes_home_override
a = anon_challenge.BrowserChallenge("https://portal.example.test/challenge?code=a", True, 600, 2, "check")
b = anon_challenge.BrowserChallenge("https://portal.example.test/challenge?code=b", False, 600, 2, "check")
anon_challenge._set_pending(a)
anon_challenge._record(a, new_attempt=False)
scope = set_hermes_home_override(tmp_path / "secondary")
try:
anon_challenge._set_pending(b)
anon_challenge._record(b, new_attempt=False)
assert not anon_challenge.record_host_outcome(a.url, 0, "done")
assert anon_challenge.pending_challenge()["url"] == b.url
finally: