diff --git a/backend/apps/subscription/free_trial.py b/backend/apps/subscription/free_trial.py index 089515e2..de7170ac 100644 --- a/backend/apps/subscription/free_trial.py +++ b/backend/apps/subscription/free_trial.py @@ -139,6 +139,13 @@ async def clear_free_trial(settings_obj) -> None: (so the UI knows it's spent) and never touches a real paid mode.""" if getattr(settings_obj, "connection_mode", "own_key") == "free-trial": settings_obj.connection_mode = "own_key" + # arm() pinned default_model to "haiku" for the free run; once the wheel is + # handed back, don't let that forced pick linger (it'd silently default a + # real subscription user to Haiku). "sonnet" is the fresh default; the + # frontend's DefaultModelGuard reconciles it to a reachable model if the + # connected provider isn't Anthropic. + if getattr(settings_obj, "default_model", None) == "haiku": + settings_obj.default_model = "sonnet" settings_obj.free_trial_token = None await save_settings_async(settings_obj) await _sync_routing(settings_obj) @@ -152,7 +159,22 @@ async def arm_free_trial(settings_obj) -> dict: mode = getattr(settings_obj, "connection_mode", "own_key") if mode not in ("own_key", "free-trial"): return {"armed": False, "reason": "other_mode"} - if _has_own_model(settings_obj) or await _has_connected_subscription(): + own = _has_own_model(settings_obj) + 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 + # up yet. Without this wait _has_connected_subscription() reads False and + # we'd arm the free trial OVER a real Claude/ChatGPT/Gemini sub, pinning the + # user to Haiku until they manually reload. Bring 9Router up so the sub is + # actually visible before we decide. Bounded + idempotent (shares the start + # lock with the boot auto-start), and skipped when a settings-level model + # already proves there's nothing to shadow. + try: + from backend.apps.nine_router import ensure_running as _ensure_9r + await _ensure_9r() + except Exception: + pass + if own or await _has_connected_subscription(): # 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 ac61df64..2e7b2df2 100644 --- a/backend/tests/test_free_trial.py +++ b/backend/tests/test_free_trial.py @@ -2,6 +2,8 @@ import backend # noqa: F401 (path sanity asserted below) +import pytest + from backend.apps.settings.models import AppSettings from backend.apps.settings.credentials import proxy_auth from backend.apps.agents.core.error_classify import ( @@ -9,7 +11,8 @@ from backend.apps.agents.core.error_classify import ( _is_transient_capacity_error, ) from backend.apps.agents.providers.registry import resolve_model_id_for_sdk -from backend.apps.subscription.free_trial import _has_own_model +from backend.apps.subscription import free_trial as ft +from backend.apps.subscription.free_trial import _has_own_model, arm_free_trial, clear_free_trial def test_proxy_auth_for_each_mode(): @@ -55,3 +58,61 @@ def test_has_own_model_never_shadows_a_real_provider(): assert _has_own_model( AppSettings(connection_mode="openswarm-pro", openswarm_bearer_token="b") ) + + +@pytest.mark.asyncio +async def test_arm_waits_for_9router_before_shadowing_a_background_started_sub(monkeypatch): + """The regression: 9Router starts in the background, so at first-boot mint time + a real Claude sub is invisible. arm() must bring 9Router up (so the sub becomes + visible) BEFORE deciding, instead of arming the free trial over it.""" + saved: list = [] + monkeypatch.setattr(ft, "save_settings_async", _record(saved)) + monkeypatch.setattr(ft, "_sync_routing", _noop) + + started = {"called": False} + + async def fake_ensure_running(): + started["called"] = True # 9Router comes up here; the sub is now visible + + # The sub is only reachable AFTER ensure_running ran (mirrors the real race). + async def sub_visible_after_start(): + return started["called"] + + import backend.apps.nine_router as nr + monkeypatch.setattr(nr, "ensure_running", fake_ensure_running) + monkeypatch.setattr(ft, "_has_connected_subscription", sub_visible_after_start) + + s = AppSettings() # no key, own_key mode: a subscription-only user + out = await arm_free_trial(s) + + assert started["called"], "arm must start 9Router before trusting the sub check" + assert out["armed"] is False and out["reason"] == "has_model" + assert s.connection_mode == "own_key" + assert s.default_model != "haiku" + + +@pytest.mark.asyncio +async def test_clear_reverts_forced_haiku_so_it_doesnt_outlive_the_trial(monkeypatch): + monkeypatch.setattr(ft, "save_settings_async", _noop) + monkeypatch.setattr(ft, "_sync_routing", _noop) + + s = AppSettings(connection_mode="free-trial", free_trial_token="ftk", default_model="haiku") + await clear_free_trial(s) + assert s.connection_mode == "own_key" + assert s.default_model == "sonnet" # forced free-run pick handed back, not left on Haiku + assert s.free_trial_token is None + + # A user who deliberately picked haiku OUTSIDE free-trial mode is left alone. + s2 = AppSettings(connection_mode="own_key", default_model="haiku") + await clear_free_trial(s2) + assert s2.default_model == "haiku" + + +async def _noop(*_a, **_k): + return None + + +def _record(bucket): + async def _inner(obj, *_a, **_k): + bucket.append(obj) + return _inner