From c752a516c0d70e172bb5d3bce1573615344d091a Mon Sep 17 00:00:00 2001 From: Petr Lavrov Date: Sat, 19 Sep 2026 09:12:13 +0300 Subject: [PATCH 1/2] feat(subscription_manager): per-user trial daily caps and no-expiry mode Public free tier: trial_duration_days=0 means no expiry. Optional per-user UTC daily caps (0 = off) for audio minutes/requests and chat requests. Hitting a daily cap sets Decision.reason=trial_daily_cap and an i18n string with minutes left today and the UTC reset time. Co-Authored-By: Claude Fable 5.1 --- .../components/middlewares/i18n_strings.py | 10 ++ botspot/components/new/DEV_README.md | 8 +- .../new/subscription_manager/__init__.py | 6 + .../new/subscription_manager/manager.py | 22 ++- .../new/subscription_manager/models.py | 1 + .../new/subscription_manager/settings.py | 13 +- .../new/subscription_manager/trial.py | 115 +++++++++++++++- example.env | 4 + .../subscription_manager_demo/README.md | 7 + .../components/new/test_component_env_vars.py | 4 + tests/components/test_subscription_manager.py | 130 ++++++++++++++++++ 11 files changed, 309 insertions(+), 11 deletions(-) diff --git a/botspot/components/middlewares/i18n_strings.py b/botspot/components/middlewares/i18n_strings.py index 626347a..a6b8be1 100644 --- a/botspot/components/middlewares/i18n_strings.py +++ b/botspot/components/middlewares/i18n_strings.py @@ -378,6 +378,16 @@ "en": "Global free trial budget reached for today (hard cost cap). Try tomorrow or add credits in /account.", "ru": "Глобальный бюджет бесплатного периода на сегодня достигнут (лимит затрат). Попробуйте завтра или пополните баланс через /account.", }, + "trial_daily_cap": { + "en": ( + "Daily free limit reached. You have {minutes_left} minutes left today. " + "Resets at {resets_at}." + ), + "ru": ( + "Дневной бесплатный лимит исчерпан. Сегодня осталось {minutes_left} мин. " + "Сброс в {resets_at}." + ), + }, "subscription_picker": { "en": "Choose a plan. Pay with Telegram Stars. Credit packs stay available as pay-as-you-go.", "ru": "Выберите план. Оплата — Telegram Stars. Пакеты кредитов остаются как оплата по факту.", diff --git a/botspot/components/new/DEV_README.md b/botspot/components/new/DEV_README.md index f803c61..4a174af 100644 --- a/botspot/components/new/DEV_README.md +++ b/botspot/components/new/DEV_README.md @@ -6,4 +6,10 @@ chat_binder.py -> components/features/chat_binder.py subscription_manager stays in `components/new/` (package). Public re-export: `botspot/subscription_manager.py`. Enable with `BOTSPOT_SUBSCRIPTION_MANAGER_ENABLED` -and MongoDB. Friends/admins bypass via `is_friend` / `is_admin`. \ No newline at end of file +and MongoDB. Friends/admins bypass via `is_friend` / `is_admin`. + +Per-user daily trial caps (0 = off): +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_MINUTES_PER_DAY`, +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_REQUESTS_PER_DAY`, +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_CHAT_REQUESTS_PER_DAY`. +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_DURATION_DAYS=0` means no expiry. \ No newline at end of file diff --git a/botspot/components/new/subscription_manager/__init__.py b/botspot/components/new/subscription_manager/__init__.py index 15333ad..fb73615 100644 --- a/botspot/components/new/subscription_manager/__init__.py +++ b/botspot/components/new/subscription_manager/__init__.py @@ -3,6 +3,12 @@ Credits ledger, Telegram Stars one-off invoices and native subscriptions, Stripe/YooKassa/TON adapters (Stripe customers + Customer Portal), trial caps, and friend/admin bypass. + +Per-user daily trial caps (env prefix ``BOTSPOT_SUBSCRIPTION_MANAGER_``, 0 = off): +``TRIAL_USER_AUDIO_MINUTES_PER_DAY``, ``TRIAL_USER_AUDIO_REQUESTS_PER_DAY``, +``TRIAL_USER_CHAT_REQUESTS_PER_DAY``. ``TRIAL_DURATION_DAYS=0`` means no expiry +(public free tier). Hitting a per-user daily cap sets ``Decision.reason`` to +``trial_daily_cap``. """ from __future__ import annotations diff --git a/botspot/components/new/subscription_manager/manager.py b/botspot/components/new/subscription_manager/manager.py index 81d5276..bfe1b16 100644 --- a/botspot/components/new/subscription_manager/manager.py +++ b/botspot/components/new/subscription_manager/manager.py @@ -49,7 +49,10 @@ from botspot.components.new.subscription_manager import stripe as stripe_adapter from botspot.components.new.subscription_manager import ton as ton_adapter from botspot.components.new.subscription_manager import yookassa as yookassa_adapter -from botspot.components.new.subscription_manager.trial import TrialLimiter +from botspot.components.new.subscription_manager.trial import ( + REASON_TRIAL_DAILY_CAP, + TrialLimiter, +) from botspot.core.errors import SubscriptionValidationError from botspot.utils.internal import get_logger @@ -439,6 +442,23 @@ async def authorize( if trial.allowed: return Decision(allowed=True, source=SOURCE_TRIAL, reason="trial") + if trial.message_key == REASON_TRIAL_DAILY_CAP: + minutes_left = ( + 0.0 if trial.minutes_left_today is None else float(trial.minutes_left_today) + ) + resets_at = trial.resets_at or "" + return Decision( + allowed=False, + source=SOURCE_NONE, + reason=REASON_TRIAL_DAILY_CAP, + message_key=REASON_TRIAL_DAILY_CAP, + user_message=t( + REASON_TRIAL_DAILY_CAP, + minutes_left=f"{minutes_left:.1f}", + resets_at=resets_at, + ), + ) + if sub and needed_minutes > 0: message_key = "subscription_minutes_exhausted" elif credits: diff --git a/botspot/components/new/subscription_manager/models.py b/botspot/components/new/subscription_manager/models.py index 2474cbe..6ce6222 100644 --- a/botspot/components/new/subscription_manager/models.py +++ b/botspot/components/new/subscription_manager/models.py @@ -82,6 +82,7 @@ class Decision(BaseModel): hold_id: str | None = None reason: str = "" message_key: str = "" + user_message: str | None = None class Balance(BaseModel): diff --git a/botspot/components/new/subscription_manager/settings.py b/botspot/components/new/subscription_manager/settings.py index e2ac9de..b33d0d8 100644 --- a/botspot/components/new/subscription_manager/settings.py +++ b/botspot/components/new/subscription_manager/settings.py @@ -73,7 +73,13 @@ class SubscriptionManagerSettings(BaseSettings): - """Env prefix ``BOTSPOT_SUBSCRIPTION_MANAGER_``.""" + """Env prefix ``BOTSPOT_SUBSCRIPTION_MANAGER_``. + + Per-user daily trial caps (0 = off): ``TRIAL_USER_AUDIO_MINUTES_PER_DAY``, + ``TRIAL_USER_AUDIO_REQUESTS_PER_DAY``, ``TRIAL_USER_CHAT_REQUESTS_PER_DAY``. + ``TRIAL_DURATION_DAYS=0`` means no expiry (public free tier); default 7 expires + after a week. + """ enabled: bool = False register_commands: bool = True @@ -88,11 +94,14 @@ class SubscriptionManagerSettings(BaseSettings): plans_json: str = "" packs_json: str = "" trial_enabled: bool = True - trial_duration_days: int = 7 + trial_duration_days: int = 7 # 0 = no expiry (public free tier) trial_user_audio_requests_total: int = 20 trial_user_audio_minutes_total: float = 90.0 trial_user_chat_requests_total: int = 200 trial_user_chat_tokens_total: int = 60000 + trial_user_audio_minutes_per_day: float = 0 # 0 = off + trial_user_audio_requests_per_day: int = 0 # 0 = off + trial_user_chat_requests_per_day: int = 0 # 0 = off trial_global_audio_requests_per_day: int = 100 trial_global_audio_minutes_per_day: float = 480.0 trial_global_chat_requests_per_day: int = 1000 diff --git a/botspot/components/new/subscription_manager/trial.py b/botspot/components/new/subscription_manager/trial.py index 321fcbb..bb4ca84 100644 --- a/botspot/components/new/subscription_manager/trial.py +++ b/botspot/components/new/subscription_manager/trial.py @@ -1,4 +1,9 @@ -"""Trial caps (whisper trial_limits.py), friends/admins bypass via user_ops.""" +"""Trial caps (whisper trial_limits.py), friends/admins bypass via user_ops. + +Per-user daily caps (0 = off) live in ``trial_daily_usage`` as one doc per user +per UTC day (``_id`` = ``{user_id}:{YYYY-MM-DD}``), incremented with atomic +``$inc``. ``trial_duration_days=0`` means no expiry. +""" from __future__ import annotations @@ -11,6 +16,8 @@ logger = get_logger() +REASON_TRIAL_DAILY_CAP = "trial_daily_cap" + @dataclass class TrialDecision: @@ -18,6 +25,8 @@ class TrialDecision: message_key: str = "" alert_message: Optional[str] = None bypassed: bool = False + minutes_left_today: Optional[float] = None + resets_at: Optional[str] = None class TrialLimiter: @@ -58,16 +67,57 @@ def _parse_alert_thresholds(raw_thresholds: str) -> list[float]: values = [0.5, 1.0, 1.5, 2.0] return sorted(set(values)) + @staticmethod + def _now_utc() -> datetime.datetime: + return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + @staticmethod def _utc_day_key(now: datetime.datetime) -> str: return now.strftime("%Y-%m-%d") + @staticmethod + def _user_day_id(user_id: int, day_key: str) -> str: + return f"{int(user_id)}:{day_key}" + + @staticmethod + def _resets_at_utc(now: datetime.datetime) -> str: + tomorrow = now.date() + datetime.timedelta(days=1) + nxt = datetime.datetime.combine(tomorrow, datetime.time.min) + return nxt.strftime("%Y-%m-%d 00:00 UTC") + @staticmethod def _limit_exceeded(current: float, increment: float, limit: float) -> bool: if limit <= 0: return False return (current + increment) > limit + def _trial_has_expiry(self) -> bool: + return int(self.settings.trial_duration_days) > 0 + + def _per_user_daily_caps_enabled(self) -> bool: + return ( + float(self.settings.trial_user_audio_minutes_per_day) > 0 + or int(self.settings.trial_user_audio_requests_per_day) > 0 + or int(self.settings.trial_user_chat_requests_per_day) > 0 + ) + + def _minutes_left_today(self, user_day: Optional[dict]) -> float: + cap = float(self.settings.trial_user_audio_minutes_per_day) + if cap <= 0: + return 0.0 + used = float((user_day or {}).get("audio_minutes_reserved", 0.0) or 0.0) + return max(0.0, cap - used) + + def _daily_cap_denied( + self, user_day: Optional[dict], now_utc: datetime.datetime + ) -> TrialDecision: + return TrialDecision( + allowed=False, + message_key=REASON_TRIAL_DAILY_CAP, + minutes_left_today=self._minutes_left_today(user_day), + resets_at=self._resets_at_utc(now_utc), + ) + def is_bypass_user(self, user_id: int) -> bool: try: from botspot.utils.user_ops import is_admin, is_friend @@ -100,15 +150,16 @@ async def reserve( chat_tokens = max(1, int(estimated_tokens)) if kind == "chat" else 0 estimated_cost_usd = max(0.0, float(estimated_cost_usd)) - now_utc = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + now_utc = self._now_utc() day_key = self._utc_day_key(now_utc) await self._upsert_trial_user_identity(user_id, username, now_utc) trial_user = await self.trial_users.find_one({"_id": user_id}) assert trial_user is not None - expires_at = trial_user.get("expires_at") - if isinstance(expires_at, datetime.datetime) and now_utc > expires_at: - return TrialDecision(allowed=False, message_key="trial_expired") + if self._trial_has_expiry(): + expires_at = trial_user.get("expires_at") + if isinstance(expires_at, datetime.datetime) and now_utc > expires_at: + return TrialDecision(allowed=False, message_key="trial_expired") if self._limit_exceeded( int(trial_user.get("audio_requests_used_total", 0)), @@ -135,6 +186,36 @@ async def reserve( ): return TrialDecision(allowed=False, message_key="trial_chat_tokens_cap") + user_day = None + if self._per_user_daily_caps_enabled(): + user_day = await self.trial_daily_usage.find_one( + {"_id": self._user_day_id(user_id, day_key)} + ) + if user_day is None: + user_day = { + "audio_requests_reserved": 0, + "audio_minutes_reserved": 0.0, + "chat_requests_reserved": 0, + } + if self._limit_exceeded( + int(user_day.get("audio_requests_reserved", 0)), + audio_requests, + float(self.settings.trial_user_audio_requests_per_day), + ): + return self._daily_cap_denied(user_day, now_utc) + if self._limit_exceeded( + float(user_day.get("audio_minutes_reserved", 0.0)), + audio_minutes, + float(self.settings.trial_user_audio_minutes_per_day), + ): + return self._daily_cap_denied(user_day, now_utc) + if self._limit_exceeded( + int(user_day.get("chat_requests_reserved", 0)), + chat_requests, + float(self.settings.trial_user_chat_requests_per_day), + ): + return self._daily_cap_denied(user_day, now_utc) + daily = await self.trial_daily_usage.find_one({"_id": day_key}) if daily is None: daily = { @@ -209,6 +290,24 @@ async def reserve( }, upsert=True, ) + if self._per_user_daily_caps_enabled(): + await self.trial_daily_usage.update_one( + {"_id": self._user_day_id(user_id, day_key)}, + { + "$setOnInsert": { + "user_id": int(user_id), + "day": day_key, + "created_at": now_utc, + }, + "$inc": { + "audio_requests_reserved": audio_requests, + "audio_minutes_reserved": audio_minutes, + "chat_requests_reserved": chat_requests, + }, + "$set": {"updated_at": now_utc}, + }, + upsert=True, + ) updated_daily = await self.trial_daily_usage.find_one({"_id": day_key}) alert_message = await self._maybe_build_alert_message(updated_daily, now_utc) return TrialDecision(allowed=True, alert_message=alert_message, bypassed=False) @@ -216,7 +315,9 @@ async def reserve( async def is_trial_active(self, user_id: int) -> bool: if not self.settings.trial_enabled: return False - now_utc = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + if not self._trial_has_expiry(): + return True + now_utc = self._now_utc() trial_user = await self.trial_users.find_one({"_id": int(user_id)}) if trial_user is None: return True @@ -229,7 +330,7 @@ async def _upsert_trial_user_identity( self, user_id: int, username: Optional[str], now_utc: datetime.datetime ) -> None: trial_days = int(self.settings.trial_duration_days) - expires_at = now_utc + datetime.timedelta(days=trial_days) + expires_at = now_utc + datetime.timedelta(days=trial_days) if trial_days > 0 else None await self.trial_users.update_one( {"_id": user_id}, { diff --git a/example.env b/example.env index 31d501f..aa5c1be 100644 --- a/example.env +++ b/example.env @@ -138,6 +138,10 @@ BOTSPOT_SEND_SAFE_ENABLED=true #BOTSPOT_SUBSCRIPTION_MANAGER_TON_USD_RATE=0 #BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_ENABLED=true #BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_DURATION_DAYS=7 +# 0 = no expiry (public free tier). Per-user daily caps: 0 = off. +#BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_MINUTES_PER_DAY=0 +#BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_REQUESTS_PER_DAY=0 +#BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_CHAT_REQUESTS_PER_DAY=0 #BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_GLOBAL_COST_USD_PER_DAY=2.0 # ---------------------------------------- diff --git a/examples/components_examples/subscription_manager_demo/README.md b/examples/components_examples/subscription_manager_demo/README.md index 1916d7e..37b9120 100644 --- a/examples/components_examples/subscription_manager_demo/README.md +++ b/examples/components_examples/subscription_manager_demo/README.md @@ -6,3 +6,10 @@ bypass billing (`BOTSPOT_FRIENDS_STR` / `BOTSPOT_ADMINS_STR`). Commands registered by the component: `/subscribe`, `/plans`, `/account`, `/buy`, `/manage` (Stripe Customer Portal), admin `/grant`, `/revoke`, `/subscribers`, `/grant_credits`. + +Per-user daily trial caps (0 = off): +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_MINUTES_PER_DAY`, +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_REQUESTS_PER_DAY`, +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_CHAT_REQUESTS_PER_DAY`. +`BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_DURATION_DAYS=0` means no expiry (public free +tier). Hitting a daily cap returns `Decision.reason=trial_daily_cap`. diff --git a/tests/components/new/test_component_env_vars.py b/tests/components/new/test_component_env_vars.py index 9b433af..55ba523 100644 --- a/tests/components/new/test_component_env_vars.py +++ b/tests/components/new/test_component_env_vars.py @@ -78,10 +78,14 @@ def test_subscription_manager_env_vars(): mp.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_ENABLED", "True") mp.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_GRACE_DAYS", "3") mp.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_STARS_PER_USD", "80") + mp.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_MINUTES_PER_DAY", "20") + mp.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_DURATION_DAYS", "0") settings = SubscriptionManagerSettings() assert settings.enabled is True assert settings.grace_days == 3 assert settings.stars_per_usd == 80 + assert settings.trial_user_audio_minutes_per_day == 20 + assert settings.trial_duration_days == 0 def test_error_handler_env_vars(): diff --git a/tests/components/test_subscription_manager.py b/tests/components/test_subscription_manager.py index 6ad042e..9dd4b5f 100644 --- a/tests/components/test_subscription_manager.py +++ b/tests/components/test_subscription_manager.py @@ -1128,3 +1128,133 @@ async def test_stripe_invoice_paid_replay_still_processes(monkeypatch, manager): assert await manager.handle_stripe_event(event) is True sub = await manager._active_subscription_doc(76) assert sub["current_period_end"] > manager._now() + timedelta(days=20) + + +def test_trial_daily_cap_defaults_off(): + settings = SubscriptionManagerSettings() + assert settings.trial_user_audio_minutes_per_day == 0 + assert settings.trial_user_audio_requests_per_day == 0 + assert settings.trial_user_chat_requests_per_day == 0 + assert settings.trial_duration_days == 7 + + +def test_trial_daily_cap_settings_env(monkeypatch): + monkeypatch.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_MINUTES_PER_DAY", "20") + monkeypatch.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_AUDIO_REQUESTS_PER_DAY", "5") + monkeypatch.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_USER_CHAT_REQUESTS_PER_DAY", "10") + monkeypatch.setenv("BOTSPOT_SUBSCRIPTION_MANAGER_TRIAL_DURATION_DAYS", "0") + settings = SubscriptionManagerSettings() + assert settings.trial_user_audio_minutes_per_day == 20 + assert settings.trial_user_audio_requests_per_day == 5 + assert settings.trial_user_chat_requests_per_day == 10 + assert settings.trial_duration_days == 0 + + +def _unlimited_trial_settings(**overrides): + data = { + "trial_duration_days": 0, + "trial_user_audio_minutes_total": 0, + "trial_user_audio_requests_total": 0, + "trial_user_chat_requests_total": 0, + "trial_user_chat_tokens_total": 0, + "trial_global_audio_minutes_per_day": 0, + "trial_global_audio_requests_per_day": 0, + "trial_global_chat_requests_per_day": 0, + "trial_global_chat_tokens_per_day": 0, + "trial_global_cost_usd_per_day": 0, + } + data.update(overrides) + return _settings(**data) + + +@pytest.mark.asyncio +async def test_per_user_daily_cap_across_day_boundary(monkeypatch): + monkeypatch.setattr("botspot.utils.user_ops.is_admin", lambda user: False) + monkeypatch.setattr("botspot.utils.user_ops.is_friend", lambda user: int(user) == 2) + from botspot.components.middlewares.i18n import t + from botspot.components.new.subscription_manager.trial import TrialLimiter + + now_box = [datetime(2026, 9, 19, 15, 0, 0)] + monkeypatch.setattr(TrialLimiter, "_now_utc", lambda *a, **k: now_box[0]) + mgr = _manager(settings=_unlimited_trial_settings(trial_user_audio_minutes_per_day=20.0)) + + first = await mgr.authorize(100, kind="audio", estimated_minutes=20, estimated_cost_usd=0.01) + assert first.allowed and first.source == SOURCE_TRIAL + denied = await mgr.authorize(100, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01) + assert denied.allowed is False + assert denied.reason == "trial_daily_cap" + assert denied.message_key == "trial_daily_cap" + assert denied.user_message is not None + assert "0.0" in denied.user_message + assert "2026-09-20 00:00 UTC" in denied.user_message + rendered = t("trial_daily_cap", minutes_left="0.0", resets_at="2026-09-20 00:00 UTC") + assert denied.user_message == rendered + + other = await mgr.authorize(101, kind="audio", estimated_minutes=5, estimated_cost_usd=0.01) + assert other.allowed and other.source == SOURCE_TRIAL + + friend = await mgr.authorize(2, kind="audio", estimated_minutes=50, estimated_cost_usd=1.0) + assert friend.allowed and friend.source == SOURCE_FRIEND + + now_box[0] = datetime(2026, 9, 20, 0, 0, 1) + next_day = await mgr.authorize(100, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01) + assert next_day.allowed and next_day.source == SOURCE_TRIAL + + +@pytest.mark.asyncio +async def test_per_user_daily_request_caps(monkeypatch): + monkeypatch.setattr("botspot.utils.user_ops.is_admin", lambda user: False) + monkeypatch.setattr("botspot.utils.user_ops.is_friend", lambda user: False) + from botspot.components.new.subscription_manager.trial import TrialLimiter + + now_box = [datetime(2026, 9, 19, 12, 0, 0)] + monkeypatch.setattr(TrialLimiter, "_now_utc", lambda *a, **k: now_box[0]) + audio_mgr = _manager(settings=_unlimited_trial_settings(trial_user_audio_requests_per_day=1)) + first = await audio_mgr.authorize( + 310, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01 + ) + second = await audio_mgr.authorize( + 310, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01 + ) + assert first.allowed + assert second.allowed is False + assert second.reason == "trial_daily_cap" + + chat_mgr = _manager(settings=_unlimited_trial_settings(trial_user_chat_requests_per_day=2)) + a = await chat_mgr.authorize(311, kind="chat", estimated_tokens=8, estimated_cost_usd=0.01) + b = await chat_mgr.authorize(311, kind="chat", estimated_tokens=8, estimated_cost_usd=0.01) + c = await chat_mgr.authorize(311, kind="chat", estimated_tokens=8, estimated_cost_usd=0.01) + assert a.allowed and b.allowed + assert c.allowed is False + assert c.reason == "trial_daily_cap" + + +@pytest.mark.asyncio +async def test_trial_duration_zero_means_no_expiry(monkeypatch): + monkeypatch.setattr("botspot.utils.user_ops.is_admin", lambda user: False) + monkeypatch.setattr("botspot.utils.user_ops.is_friend", lambda user: False) + from botspot.components.new.subscription_manager.trial import TrialLimiter + + now_box = [datetime(2026, 9, 1, 12, 0, 0)] + monkeypatch.setattr(TrialLimiter, "_now_utc", lambda *a, **k: now_box[0]) + free = _manager(settings=_unlimited_trial_settings(trial_duration_days=0)) + first = await free.authorize(200, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01) + assert first.allowed + assert await free.trial.is_trial_active(200) + now_box[0] = datetime(2026, 9, 20, 12, 0, 0) + later = await free.authorize(200, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01) + assert later.allowed + assert await free.trial.is_trial_active(200) + doc = await free.trial.trial_users.find_one({"_id": 200}) + assert doc is not None + assert doc.get("expires_at") is None + + now_box[0] = datetime(2026, 9, 1, 12, 0, 0) + week = _manager(settings=_settings(trial_duration_days=7, trial_global_cost_usd_per_day=10)) + started = await week.authorize(201, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01) + assert started.allowed + now_box[0] = datetime(2026, 9, 9, 12, 0, 0) + expired = await week.authorize(201, kind="audio", estimated_minutes=1, estimated_cost_usd=0.01) + assert expired.allowed is False + assert expired.message_key == "trial_expired" + assert await week.trial.is_trial_active(201) is False From b873fb80f4949eacc5097e9599b52ed8c7170aaf Mon Sep 17 00:00:00 2001 From: Petr Lavrov Date: Sat, 19 Sep 2026 09:15:59 +0300 Subject: [PATCH 2/2] refactor(subscription_manager): split trial cap checks under CCN 20 Extract lifetime/global/per-user daily cap helpers so reserve and authorize stay under the advisory lizard limit. Co-Authored-By: Claude Fable 5.1 --- .../new/subscription_manager/manager.py | 3 +- .../new/subscription_manager/trial.py | 329 ++++++++++++------ 2 files changed, 215 insertions(+), 117 deletions(-) diff --git a/botspot/components/new/subscription_manager/manager.py b/botspot/components/new/subscription_manager/manager.py index bfe1b16..7bbeaa7 100644 --- a/botspot/components/new/subscription_manager/manager.py +++ b/botspot/components/new/subscription_manager/manager.py @@ -441,7 +441,9 @@ async def authorize( ) if trial.allowed: return Decision(allowed=True, source=SOURCE_TRIAL, reason="trial") + return self._deny_authorize(trial, sub, credits, needed_minutes) + def _deny_authorize(self, trial, sub, credits: int, needed_minutes: float) -> Decision: if trial.message_key == REASON_TRIAL_DAILY_CAP: minutes_left = ( 0.0 if trial.minutes_left_today is None else float(trial.minutes_left_today) @@ -458,7 +460,6 @@ async def authorize( resets_at=resets_at, ), ) - if sub and needed_minutes > 0: message_key = "subscription_minutes_exhausted" elif credits: diff --git a/botspot/components/new/subscription_manager/trial.py b/botspot/components/new/subscription_manager/trial.py index bb4ca84..ee0573f 100644 --- a/botspot/components/new/subscription_manager/trial.py +++ b/botspot/components/new/subscription_manager/trial.py @@ -128,97 +128,145 @@ def is_bypass_user(self, user_id: int) -> bool: logger.warning(f"Failed to check bypass role for user {user_id}: {e}") return False - async def reserve( - self, - user_id: int, - *, - kind: str, - estimated_minutes: float = 0, - estimated_tokens: int = 0, - estimated_cost_usd: float = 0, - username: Optional[str] = None, - ) -> TrialDecision: - if not self.settings.trial_enabled: - return TrialDecision(allowed=True, bypassed=True) - - if self.is_bypass_user(user_id): - return TrialDecision(allowed=True, bypassed=True) - - audio_requests = 1 if kind == "audio" else 0 - audio_minutes = max(1.0, float(estimated_minutes)) if kind == "audio" else 0.0 - chat_requests = 1 if kind == "chat" else 0 - chat_tokens = max(1, int(estimated_tokens)) if kind == "chat" else 0 - estimated_cost_usd = max(0.0, float(estimated_cost_usd)) + def _first_cap_key(self, checks: list[tuple[float, float, float, str]]) -> Optional[str]: + for current, increment, limit, message_key in checks: + if self._limit_exceeded(current, increment, limit): + return message_key + return None - now_utc = self._now_utc() - day_key = self._utc_day_key(now_utc) - await self._upsert_trial_user_identity(user_id, username, now_utc) - trial_user = await self.trial_users.find_one({"_id": user_id}) - assert trial_user is not None + def _lifetime_cap_key( + self, + trial_user: dict, + audio_requests: int, + audio_minutes: float, + chat_requests: int, + chat_tokens: int, + ) -> Optional[str]: + s = self.settings + return self._first_cap_key( + [ + ( + int(trial_user.get("audio_requests_used_total", 0)), + audio_requests, + float(s.trial_user_audio_requests_total), + "trial_audio_requests_cap", + ), + ( + float(trial_user.get("audio_minutes_used_total", 0.0)), + audio_minutes, + float(s.trial_user_audio_minutes_total), + "trial_audio_minutes_cap", + ), + ( + int(trial_user.get("chat_requests_used_total", 0)), + chat_requests, + float(s.trial_user_chat_requests_total), + "trial_chat_requests_cap", + ), + ( + int(trial_user.get("chat_tokens_used_total", 0)), + float(chat_tokens), + float(s.trial_user_chat_tokens_total), + "trial_chat_tokens_cap", + ), + ] + ) - if self._trial_has_expiry(): - expires_at = trial_user.get("expires_at") - if isinstance(expires_at, datetime.datetime) and now_utc > expires_at: - return TrialDecision(allowed=False, message_key="trial_expired") + def _global_cap_key( + self, + daily: dict, + audio_requests: int, + audio_minutes: float, + chat_requests: int, + chat_tokens: int, + estimated_cost_usd: float, + ) -> Optional[str]: + s = self.settings + return self._first_cap_key( + [ + ( + int(daily.get("audio_requests_reserved", 0)), + audio_requests, + float(s.trial_global_audio_requests_per_day), + "trial_global_audio_requests_day", + ), + ( + float(daily.get("audio_minutes_reserved", 0.0)), + audio_minutes, + float(s.trial_global_audio_minutes_per_day), + "trial_global_audio_minutes_day", + ), + ( + int(daily.get("chat_requests_reserved", 0)), + chat_requests, + float(s.trial_global_chat_requests_per_day), + "trial_global_chat_requests_day", + ), + ( + int(daily.get("chat_tokens_reserved", 0)), + float(chat_tokens), + float(s.trial_global_chat_tokens_per_day), + "trial_global_chat_tokens_day", + ), + ( + float(daily.get("cost_reserved_usd", 0.0)), + estimated_cost_usd, + float(s.trial_global_cost_usd_per_day), + "trial_global_cost_day", + ), + ] + ) - if self._limit_exceeded( - int(trial_user.get("audio_requests_used_total", 0)), - audio_requests, - float(self.settings.trial_user_audio_requests_total), - ): - return TrialDecision(allowed=False, message_key="trial_audio_requests_cap") - if self._limit_exceeded( - float(trial_user.get("audio_minutes_used_total", 0.0)), - audio_minutes, - float(self.settings.trial_user_audio_minutes_total), - ): - return TrialDecision(allowed=False, message_key="trial_audio_minutes_cap") - if self._limit_exceeded( - int(trial_user.get("chat_requests_used_total", 0)), - chat_requests, - float(self.settings.trial_user_chat_requests_total), - ): - return TrialDecision(allowed=False, message_key="trial_chat_requests_cap") - if self._limit_exceeded( - int(trial_user.get("chat_tokens_used_total", 0)), - float(chat_tokens), - float(self.settings.trial_user_chat_tokens_total), - ): - return TrialDecision(allowed=False, message_key="trial_chat_tokens_cap") - - user_day = None - if self._per_user_daily_caps_enabled(): - user_day = await self.trial_daily_usage.find_one( - {"_id": self._user_day_id(user_id, day_key)} + def _per_user_daily_hit( + self, + user_day: dict, + audio_requests: int, + audio_minutes: float, + chat_requests: int, + ) -> bool: + s = self.settings + return ( + self._first_cap_key( + [ + ( + int(user_day.get("audio_requests_reserved", 0)), + audio_requests, + float(s.trial_user_audio_requests_per_day), + REASON_TRIAL_DAILY_CAP, + ), + ( + float(user_day.get("audio_minutes_reserved", 0.0)), + audio_minutes, + float(s.trial_user_audio_minutes_per_day), + REASON_TRIAL_DAILY_CAP, + ), + ( + int(user_day.get("chat_requests_reserved", 0)), + chat_requests, + float(s.trial_user_chat_requests_per_day), + REASON_TRIAL_DAILY_CAP, + ), + ] ) - if user_day is None: - user_day = { - "audio_requests_reserved": 0, - "audio_minutes_reserved": 0.0, - "chat_requests_reserved": 0, - } - if self._limit_exceeded( - int(user_day.get("audio_requests_reserved", 0)), - audio_requests, - float(self.settings.trial_user_audio_requests_per_day), - ): - return self._daily_cap_denied(user_day, now_utc) - if self._limit_exceeded( - float(user_day.get("audio_minutes_reserved", 0.0)), - audio_minutes, - float(self.settings.trial_user_audio_minutes_per_day), - ): - return self._daily_cap_denied(user_day, now_utc) - if self._limit_exceeded( - int(user_day.get("chat_requests_reserved", 0)), - chat_requests, - float(self.settings.trial_user_chat_requests_per_day), - ): - return self._daily_cap_denied(user_day, now_utc) + is not None + ) + async def _load_user_day(self, user_id: int, day_key: str) -> dict: + user_day = await self.trial_daily_usage.find_one( + {"_id": self._user_day_id(user_id, day_key)} + ) + if user_day is None: + return { + "audio_requests_reserved": 0, + "audio_minutes_reserved": 0.0, + "chat_requests_reserved": 0, + } + return user_day + + async def _load_global_day(self, day_key: str) -> dict: daily = await self.trial_daily_usage.find_one({"_id": day_key}) if daily is None: - daily = { + return { "_id": day_key, "audio_requests_reserved": 0, "audio_minutes_reserved": 0.0, @@ -227,38 +275,21 @@ async def reserve( "cost_reserved_usd": 0.0, "alert_level": 0, } + return daily - if self._limit_exceeded( - int(daily.get("audio_requests_reserved", 0)), - audio_requests, - float(self.settings.trial_global_audio_requests_per_day), - ): - return TrialDecision(allowed=False, message_key="trial_global_audio_requests_day") - if self._limit_exceeded( - float(daily.get("audio_minutes_reserved", 0.0)), - audio_minutes, - float(self.settings.trial_global_audio_minutes_per_day), - ): - return TrialDecision(allowed=False, message_key="trial_global_audio_minutes_day") - if self._limit_exceeded( - int(daily.get("chat_requests_reserved", 0)), - chat_requests, - float(self.settings.trial_global_chat_requests_per_day), - ): - return TrialDecision(allowed=False, message_key="trial_global_chat_requests_day") - if self._limit_exceeded( - int(daily.get("chat_tokens_reserved", 0)), - float(chat_tokens), - float(self.settings.trial_global_chat_tokens_per_day), - ): - return TrialDecision(allowed=False, message_key="trial_global_chat_tokens_day") - if self._limit_exceeded( - float(daily.get("cost_reserved_usd", 0.0)), - estimated_cost_usd, - float(self.settings.trial_global_cost_usd_per_day), - ): - return TrialDecision(allowed=False, message_key="trial_global_cost_day") - + async def _record_reservation( + self, + user_id: int, + username: Optional[str], + now_utc: datetime.datetime, + day_key: str, + daily: dict, + audio_requests: int, + audio_minutes: float, + chat_requests: int, + chat_tokens: int, + estimated_cost_usd: float, + ) -> Optional[str]: await self.trial_users.update_one( {"_id": user_id}, { @@ -309,7 +340,73 @@ async def reserve( upsert=True, ) updated_daily = await self.trial_daily_usage.find_one({"_id": day_key}) - alert_message = await self._maybe_build_alert_message(updated_daily, now_utc) + return await self._maybe_build_alert_message(updated_daily, now_utc) + + async def reserve( + self, + user_id: int, + *, + kind: str, + estimated_minutes: float = 0, + estimated_tokens: int = 0, + estimated_cost_usd: float = 0, + username: Optional[str] = None, + ) -> TrialDecision: + if not self.settings.trial_enabled or self.is_bypass_user(user_id): + return TrialDecision(allowed=True, bypassed=True) + + audio_requests = 1 if kind == "audio" else 0 + audio_minutes = max(1.0, float(estimated_minutes)) if kind == "audio" else 0.0 + chat_requests = 1 if kind == "chat" else 0 + chat_tokens = max(1, int(estimated_tokens)) if kind == "chat" else 0 + estimated_cost_usd = max(0.0, float(estimated_cost_usd)) + + now_utc = self._now_utc() + day_key = self._utc_day_key(now_utc) + await self._upsert_trial_user_identity(user_id, username, now_utc) + trial_user = await self.trial_users.find_one({"_id": user_id}) + assert trial_user is not None + + if self._trial_has_expiry(): + expires_at = trial_user.get("expires_at") + if isinstance(expires_at, datetime.datetime) and now_utc > expires_at: + return TrialDecision(allowed=False, message_key="trial_expired") + + lifetime_key = self._lifetime_cap_key( + trial_user, audio_requests, audio_minutes, chat_requests, chat_tokens + ) + if lifetime_key: + return TrialDecision(allowed=False, message_key=lifetime_key) + + if self._per_user_daily_caps_enabled(): + user_day = await self._load_user_day(user_id, day_key) + if self._per_user_daily_hit(user_day, audio_requests, audio_minutes, chat_requests): + return self._daily_cap_denied(user_day, now_utc) + + daily = await self._load_global_day(day_key) + global_key = self._global_cap_key( + daily, + audio_requests, + audio_minutes, + chat_requests, + chat_tokens, + estimated_cost_usd, + ) + if global_key: + return TrialDecision(allowed=False, message_key=global_key) + + alert_message = await self._record_reservation( + user_id, + username, + now_utc, + day_key, + daily, + audio_requests, + audio_minutes, + chat_requests, + chat_tokens, + estimated_cost_usd, + ) return TrialDecision(allowed=True, alert_message=alert_message, bypassed=False) async def is_trial_active(self, user_id: int) -> bool: