diff --git a/backend/apps/subscription/free_trial.py b/backend/apps/subscription/free_trial.py index de7170ac..4f8786f1 100644 --- a/backend/apps/subscription/free_trial.py +++ b/backend/apps/subscription/free_trial.py @@ -10,6 +10,7 @@ forced cheap model; this module only mirrors state into settings and 9Router. from __future__ import annotations +import asyncio import hashlib import logging import os @@ -160,6 +161,7 @@ async def arm_free_trial(settings_obj) -> dict: if mode not in ("own_key", "free-trial"): return {"armed": False, "reason": "other_mode"} own = _has_own_model(settings_obj) + has_sub = False if not own: # A subscription lives in 9Router, not settings, and 9Router now starts in # the BACKGROUND (non-blocking boot), so at first-launch mint time it isn't @@ -174,7 +176,19 @@ async def arm_free_trial(settings_obj) -> dict: await _ensure_9r() except Exception: pass - if own or await _has_connected_subscription(): + # 9Router's /api/providers can lag /v1/models (what is_running probes) by a + # beat on a cold start, so a real sub can read as absent for a sub-second + # window. Re-check a few times before concluding "no sub", so we never arm + # over a sub that's merely still loading. CAPPED on purpose: a genuinely + # sub-less user exhausts these in ~1.2s and falls through to arm, so this + # never waits on a subscription that doesn't exist. + for _i in range(5): + if await _has_connected_subscription(): + has_sub = True + break + if _i < 4: + await asyncio.sleep(0.3) + if own or has_sub: # A real model exists now (key, custom provider, or a 9Router sub). If we # were on the free lane, hand the wheel back instead of re-arming. if mode == "free-trial": diff --git a/backend/tests/test_free_trial.py b/backend/tests/test_free_trial.py index 2e7b2df2..bc4cc522 100644 --- a/backend/tests/test_free_trial.py +++ b/backend/tests/test_free_trial.py @@ -91,6 +91,57 @@ async def test_arm_waits_for_9router_before_shadowing_a_background_started_sub(m assert s.default_model != "haiku" +@pytest.mark.asyncio +async def test_arm_tolerates_provider_load_lag(monkeypatch): + """9Router's /api/providers can lag is_running on a cold start. arm must re-check + a few times so a sub that loads a beat late is still caught, not shadowed.""" + monkeypatch.setattr(ft, "save_settings_async", _noop) + monkeypatch.setattr(ft, "_sync_routing", _noop) + + async def fake_ensure_running(): + return None + + calls = {"n": 0} + async def lagging_sub(): + calls["n"] += 1 + return calls["n"] >= 3 # empty for the first two probes, then the sub appears + + import backend.apps.nine_router as nr + monkeypatch.setattr(nr, "ensure_running", fake_ensure_running) + monkeypatch.setattr(ft, "_has_connected_subscription", lagging_sub) + + s = AppSettings() + res = await ft.arm_free_trial(s) + assert res["reason"] == "has_model", res + assert s.default_model != "haiku" + assert calls["n"] >= 3, "should have re-checked past the lagging-empty probes" + + +@pytest.mark.asyncio +async def test_arm_with_no_sub_is_bounded_and_falls_through_to_arm(monkeypatch): + """The 'don't poll for something that doesn't exist' guarantee: a genuinely + sub-less user must exhaust the re-checks quickly and PROCEED to arm, never hang.""" + async def fake_ensure_running(): + return None + async def never_sub(): + return False + + import time + import backend.apps.nine_router as nr + monkeypatch.setattr(nr, "ensure_running", fake_ensure_running) + monkeypatch.setattr(ft, "_has_connected_subscription", never_sub) + # Short-circuit before the cloud mint so the test stays offline + deterministic; + # reaching this branch proves arm did NOT falsely conclude has_model. + monkeypatch.setattr(ft, "_fingerprint", lambda _s: None) + + s = AppSettings() + t = time.monotonic() + res = await ft.arm_free_trial(s) + elapsed = time.monotonic() - t + assert res["reason"] == "no_fingerprint", res # got past the sub guard to the arm path + assert elapsed < 3.0, f"re-check budget not bounded: {elapsed:.2f}s" + + @pytest.mark.asyncio async def test_clear_reverts_forced_haiku_so_it_doesnt_outlive_the_trial(monkeypatch): monkeypatch.setattr(ft, "save_settings_async", _noop)