Enforce premium quota across sessions daily (#287)

Co-authored-by: OpenAI Codex <codex@openai.com>

agent/core/session_persistence.py

@@ -175,6 +175,7 @@ class MongoSessionStore(NoopSessionStore):
         turn_count: int = 0,
         pending_approval: list[dict[str, Any]] | None = None,
         claude_counted: bool = False,
+        claude_counted_day: str | None = None,
         premium_user_billed: bool = False,
         notification_destinations: list[str] | None = None,
         auto_approval_enabled: bool = False,
@@ -207,6 +208,7 @@ class MongoSessionStore(NoopSessionStore):
                     "turn_count": turn_count,
                     "pending_approval": pending_approval or [],
                     "claude_counted": claude_counted,
+                    "claude_counted_day": claude_counted_day,
                     "premium_user_billed": premium_user_billed,
                     "notification_destinations": notification_destinations or [],
                     "auto_approval_enabled": auto_approval_enabled,
@@ -230,6 +232,7 @@ class MongoSessionStore(NoopSessionStore):
         turn_count: int = 0,
         pending_approval: list[dict[str, Any]] | None = None,
         claude_counted: bool = False,
+        claude_counted_day: str | None = None,
         premium_user_billed: bool = False,
         created_at: datetime | None = None,
         notification_destinations: list[str] | None = None,
@@ -255,6 +258,7 @@ class MongoSessionStore(NoopSessionStore):
             turn_count=turn_count,
             pending_approval=pending_approval,
             claude_counted=claude_counted,
+            claude_counted_day=claude_counted_day,
             premium_user_billed=premium_user_billed,
             notification_destinations=notification_destinations,
             auto_approval_enabled=auto_approval_enabled,

backend/routes/agent.py

@@ -207,21 +207,23 @@ async def _enforce_premium_model_quota(
 
     Runs at *message-submit* time, not session-create time — so spinning up a
     premium-model session to look around doesn't burn quota. The
-    ``claude_counted`` flag on ``AgentSession`` guards against re-counting the
-    same session; the stored field name is kept for persistence compatibility.
+    ``claude_counted_day`` flag on ``AgentSession`` guards against re-counting
+    the same session on the same day, while still counting old sessions when
+    they are used again on a later day.
 
     Subsidizes the daily allowance (free = 2 for default premium, pro = 20
     across premium models), organization-billed through the HF Router. Opus and
     GPT-5.5 are pro-only before quota is charged. Past the allowance, premium
     router models flip the session to ``premium_user_billed`` so the call bills
     the user's own HF token instead of blocking. No-ops when the model isn't
-    premium or when this session's billing has already been decided.
+    premium or when this session's billing has already been decided for today.
     """
     model_name = agent_session.session.config.model_name
     if not _is_premium_model(model_name):
         return
     _reject_model_unavailable_for_plan(model_name, user)
-    if agent_session.claude_counted:
+    quota_day = user_quotas.current_quota_day()
+    if agent_session.claude_counted and agent_session.claude_counted_day == quota_day:
         return
     user_id = user["user_id"]
     plan = user.get("plan", "free")
@@ -241,7 +243,12 @@ async def _enforce_premium_model_quota(
         # Past the subsidized allowance on a user-billable model: bill the
         # user's own HF (OAuth) token for this session instead of blocking.
         agent_session.session.premium_user_billed = True
+    else:
+        # A session that overflowed on a previous day can use today's
+        # subsidized allowance again if quota is available.
+        agent_session.session.premium_user_billed = False
     agent_session.claude_counted = True
+    agent_session.claude_counted_day = quota_day
     await session_manager.persist_session_snapshot(agent_session)
 
 
@@ -909,8 +916,8 @@ async def chat_sse(
     approvals = body.get("approvals")
 
     # Gate user-message sends against the daily premium-model quota. Approvals are
-    # continuations of an in-progress turn — the session was already charged
-    # on its first message, so we skip the gate there.
+    # continuations of an in-progress turn, so the relevant quota decision was
+    # made when that user message was submitted.
     if text is not None and not approvals:
         try:
             await _enforce_premium_model_quota(user, agent_session)

backend/session_manager.py

@@ -6,7 +6,7 @@ import logging
 import os
 import uuid
 from dataclasses import dataclass, field
-from datetime import datetime, timedelta
+from datetime import UTC, datetime, timedelta
 from pathlib import Path
 from typing import Any, Optional
 
@@ -113,9 +113,28 @@ class AgentSession:
     is_reaping: bool = False
     broadcaster: Any = None
     title: str | None = None
-    # True once this session has been counted against the user's daily premium
-    # quota. The field name is kept for persistence compatibility.
+    # True once this session has ever been counted against premium quota.
+    # claude_counted_day decides whether it has already consumed today's cap.
     claude_counted: bool = False
+    claude_counted_day: str | None = None
+
+
+def _quota_day_today() -> str:
+    return datetime.now(UTC).date().isoformat()
+
+
+def _quota_counted_today(agent_session: AgentSession) -> bool:
+    return (
+        agent_session.claude_counted
+        and agent_session.claude_counted_day == _quota_day_today()
+    )
+
+
+def _premium_user_billed_today(agent_session: AgentSession) -> bool:
+    return bool(
+        getattr(agent_session.session, "premium_user_billed", False)
+        and _quota_counted_today(agent_session)
+    )
 
 
 class SessionCapacityError(Exception):
@@ -600,6 +619,7 @@ class SessionManager:
                     agent_session.session
                 ),
                 claude_counted=agent_session.claude_counted,
+                claude_counted_day=agent_session.claude_counted_day,
                 premium_user_billed=getattr(
                     agent_session.session, "premium_user_billed", False
                 ),
@@ -696,6 +716,13 @@ class SessionManager:
             premium_user_billed=bool(meta.get("premium_user_billed", False)),
             claude_counted=bool(meta.get("claude_counted")),
         )
+        claude_counted_day = (
+            str(meta.get("claude_counted_day"))
+            if meta.get("claude_counted_day")
+            else None
+        )
+        if not claude_counted:
+            claude_counted_day = None
         event_queue: asyncio.Queue = asyncio.Queue()
         submission_queue: asyncio.Queue = asyncio.Queue()
         tool_router, session = await asyncio.to_thread(
@@ -779,6 +806,7 @@ class SessionManager:
             is_active=True,
             is_processing=False,
             claude_counted=claude_counted,
+            claude_counted_day=claude_counted_day,
             title=meta.get("title"),
         )
         started = await self._start_agent_session(
@@ -1505,10 +1533,8 @@ class SessionManager:
                 agent_session.session.notification_destinations
             ),
             "auto_approval": self._auto_approval_summary(agent_session.session),
-            "premium_user_billed": getattr(
-                agent_session.session, "premium_user_billed", False
-            ),
-            "premium_quota_counted": agent_session.claude_counted,
+            "premium_user_billed": _premium_user_billed_today(agent_session),
+            "premium_quota_counted": _quota_counted_today(agent_session),
         }
 
     def set_notification_destinations(
@@ -1562,6 +1588,10 @@ class SessionManager:
                 else:
                     created_at_str = str(created_at or datetime.utcnow().isoformat())
                 pending = self._pending_docs_for_api(row.get("pending_approval") or [])
+                quota_counted_today = (
+                    bool(row.get("claude_counted", False))
+                    and row.get("claude_counted_day") == _quota_day_today()
+                )
                 results.append(
                     {
                         "session_id": str(sid),
@@ -1575,8 +1605,9 @@ class SessionManager:
                         "title": row.get("title"),
                         "premium_user_billed": bool(
                             row.get("premium_user_billed", False)
+                            and quota_counted_today
                         ),
-                        "premium_quota_counted": bool(row.get("claude_counted", False)),
+                        "premium_quota_counted": quota_counted_today,
                         "notification_destinations": row.get(
                             "notification_destinations"
                         )

backend/user_quotas.py

@@ -7,10 +7,10 @@ in-process dict remains the fallback for local/dev/test runs.
 The public names still say ``claude`` because this quota bucket originally
 only covered Claude and the persisted session field uses that name.
 
-Unit: first premium-model submit in a session, not raw messages. A user who
-sends with an allowed premium model in a new session consumes one quota point;
-switching an already-counted session back to a premium model doesn't
-(`AgentSession.claude_counted` guards that). Model-level plan gates live in
+Unit: first premium-model submit per session per UTC day, not raw messages. A
+user who sends with an allowed premium model in any session consumes one quota
+point for that day; continuing the same session on the same day doesn't
+(`AgentSession.claude_counted_day` guards that). Model-level plan gates live in
 ``backend.routes.agent``; this module only tracks the per-plan daily cap.
 
 Cap tiers:
@@ -40,6 +40,11 @@ def _today() -> str:
     return datetime.now(UTC).date().isoformat()
 
 
+def current_quota_day() -> str:
+    """Return the UTC date key used for today's premium-model quota bucket."""
+    return _today()
+
+
 def daily_cap_for(plan: str | None) -> int:
     """Return the daily Claude-session cap for the given plan."""
     return CLAUDE_PRO_DAILY if plan == "pro" else CLAUDE_FREE_DAILY

frontend/src/components/SessionChat.tsx

@@ -94,13 +94,13 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
   // account (past the subsidized daily allowance). Re-read after each turn,
   // since the backend flips it at submit time. Only premium-model sessions can
   // ever be user-billed, so skip the fetch for free models.
-  const [premiumBilled, setPremiumBilled] = useState(false);
-  const [premiumQuotaCounted, setPremiumQuotaCounted] = useState(false);
+  const [premiumBilled, setPremiumBilled] = useState<boolean | null>(null);
+  const [premiumQuotaCounted, setPremiumQuotaCounted] = useState<boolean | null>(null);
   const onPremiumModel = isPremiumPath(sessionMeta?.model ?? undefined);
   useEffect(() => {
     if (!isActive || !onPremiumModel) {
-      setPremiumBilled(false);
-      setPremiumQuotaCounted(false);
+      setPremiumBilled(null);
+      setPremiumQuotaCounted(null);
       return;
     }
     if (busy) return;
@@ -119,9 +119,9 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
     };
   }, [busy, isActive, onPremiumModel, sessionId]);
 
-  const sessionPremiumBilled = premiumBilled || Boolean(sessionMeta?.premiumUserBilled);
+  const sessionPremiumBilled = premiumBilled ?? Boolean(sessionMeta?.premiumUserBilled);
   const sessionPremiumQuotaCounted =
-    premiumQuotaCounted || Boolean(sessionMeta?.premiumQuotaCounted);
+    premiumQuotaCounted ?? Boolean(sessionMeta?.premiumQuotaCounted);
   const premiumBillingNotice =
     sessionPremiumBilled ||
     (isActive &&

frontend/src/store/sessionStore.ts

@@ -252,10 +252,14 @@ export const useSessionStore = create<SessionStore>()(
     {
       name: 'hf-agent-sessions',
       partialize: (state) => ({
-        // Reset the transient isProcessing flag so a stale `true` from a
-        // previous session can't trigger a reactivating hydration on the next
-        // cold load — it's always re-derived from the live GET /sessions list.
-        sessions: state.sessions.map((s) => ({ ...s, isProcessing: false })),
+        // Reset transient session flags so cold-load state is re-derived from
+        // the live GET /sessions list and current daily quota window.
+        sessions: state.sessions.map((s) => ({
+          ...s,
+          isProcessing: false,
+          premiumUserBilled: false,
+          premiumQuotaCounted: false,
+        })),
         activeSessionId: state.activeSessionId,
       }),
     }

tests/unit/test_agent_model_gating.py

@@ -24,6 +24,7 @@ def _reset_quota_store():
 def _premium_session(model: str = agent.DEFAULT_PREMIUM_MODEL_ID):
     return SimpleNamespace(
         claude_counted=False,
+        claude_counted_day=None,
         session=SimpleNamespace(
             config=SimpleNamespace(model_name=model),
             premium_user_billed=False,
@@ -198,11 +199,66 @@ async def test_premium_quota_charges_without_user_billing_inside_allowance(monke
     )
 
     assert agent_session.claude_counted is True
+    assert agent_session.claude_counted_day == agent.user_quotas.current_quota_day()
     assert agent_session.session.premium_user_billed is False
     assert persisted == [agent_session]
     assert await agent.user_quotas.get_claude_used_today("u1") == 1
 
 
+@pytest.mark.asyncio
+async def test_premium_quota_counts_same_session_once_per_day(monkeypatch):
+    async def fake_persist_session_snapshot(_agent_session):
+        return None
+
+    monkeypatch.setattr(
+        agent.session_manager,
+        "persist_session_snapshot",
+        fake_persist_session_snapshot,
+    )
+
+    agent_session = _premium_session()
+
+    await agent._enforce_premium_model_quota(
+        {"user_id": "u1", "plan": "free"},
+        agent_session,
+    )
+    await agent._enforce_premium_model_quota(
+        {"user_id": "u1", "plan": "free"},
+        agent_session,
+    )
+
+    assert agent_session.claude_counted is True
+    assert agent_session.claude_counted_day == agent.user_quotas.current_quota_day()
+    assert await agent.user_quotas.get_claude_used_today("u1") == 1
+
+
+@pytest.mark.asyncio
+async def test_premium_quota_counts_stale_session_again_today(monkeypatch):
+    async def fake_persist_session_snapshot(_agent_session):
+        return None
+
+    monkeypatch.setattr(
+        agent.session_manager,
+        "persist_session_snapshot",
+        fake_persist_session_snapshot,
+    )
+
+    agent_session = _premium_session()
+    agent_session.claude_counted = True
+    agent_session.claude_counted_day = "2000-01-01"
+    agent_session.session.premium_user_billed = True
+
+    await agent._enforce_premium_model_quota(
+        {"user_id": "u1", "plan": "free"},
+        agent_session,
+    )
+
+    assert agent_session.claude_counted is True
+    assert agent_session.claude_counted_day == agent.user_quotas.current_quota_day()
+    assert agent_session.session.premium_user_billed is False
+    assert await agent.user_quotas.get_claude_used_today("u1") == 1
+
+
 @pytest.mark.asyncio
 async def test_free_user_gets_two_subsidized_premium_sessions_then_user_billing(
     monkeypatch,
@@ -227,6 +283,7 @@ async def test_free_user_gets_two_subsidized_premium_sessions_then_user_billing(
     third = _premium_session()
     await agent._enforce_premium_model_quota({"user_id": "g1", "plan": "free"}, third)
     assert third.session.premium_user_billed is True
+    assert third.claude_counted_day == agent.user_quotas.current_quota_day()
     assert await agent.user_quotas.get_claude_used_today("g1") == 2
 
 
@@ -249,6 +306,7 @@ async def test_free_model_does_not_consume_premium_quota(monkeypatch):
     )
 
     assert agent_session.claude_counted is False
+    assert agent_session.claude_counted_day is None
     assert await agent.user_quotas.get_claude_used_today("u1") == 0
 
 
@@ -274,6 +332,7 @@ async def test_free_user_cannot_spend_quota_on_pro_only_premium_model(monkeypatc
     assert exc_info.value.status_code == 403
     assert exc_info.value.detail["error"] == "model_requires_pro"
     assert agent_session.claude_counted is False
+    assert agent_session.claude_counted_day is None
     assert await agent.user_quotas.get_claude_used_today("u1") == 0
 
 
@@ -321,6 +380,7 @@ async def test_pro_user_uses_pro_premium_quota(monkeypatch):
             agent_session,
         )
         assert agent_session.claude_counted is True
+        assert agent_session.claude_counted_day == agent.user_quotas.current_quota_day()
         assert agent_session.session.premium_user_billed is False
         assert await agent.user_quotas.get_claude_used_today("pro-user") == index + 1
 

tests/unit/test_session_manager_persistence.py

@@ -757,6 +757,7 @@ async def test_list_sessions_dev_uses_store_dev_visibility():
                         "created_at": datetime.now(UTC),
                         "premium_user_billed": True,
                         "claude_counted": True,
+                        "claude_counted_day": datetime.now(UTC).date().isoformat(),
                         "auto_approval_enabled": True,
                         "auto_approval_cost_cap_usd": 5.0,
                         "auto_approval_estimated_spend_usd": 2.0,
@@ -786,3 +787,44 @@ async def test_list_sessions_dev_uses_store_dev_visibility():
         "estimated_spend_usd": 2.0,
         "remaining_usd": 3.0,
     }
+
+
+def test_get_session_info_marks_stale_premium_quota_as_unused_today():
+    manager = _manager_with_store(NoopSessionStore())
+    agent_session = _runtime_agent_session("s1", user_id="alice")
+    agent_session.claude_counted = True
+    agent_session.claude_counted_day = "2000-01-01"
+    agent_session.session.premium_user_billed = True
+    manager.sessions["s1"] = agent_session
+
+    info = manager.get_session_info("s1")
+
+    assert info is not None
+    assert info["premium_user_billed"] is False
+    assert info["premium_quota_counted"] is False
+
+
+@pytest.mark.asyncio
+async def test_list_sessions_marks_stale_premium_quota_as_unused_today():
+    class ListStore(NoopSessionStore):
+        enabled = True
+
+        async def list_sessions(self, user_id: str, **_: Any) -> list[dict[str, Any]]:
+            return [
+                {
+                    "session_id": "s1",
+                    "user_id": user_id,
+                    "model": "m",
+                    "created_at": datetime.now(UTC),
+                    "premium_user_billed": True,
+                    "claude_counted": True,
+                    "claude_counted_day": "2000-01-01",
+                }
+            ]
+
+    manager = _manager_with_store(ListStore())
+
+    sessions = await manager.list_sessions(user_id="alice")
+
+    assert sessions[0]["premium_user_billed"] is False
+    assert sessions[0]["premium_quota_counted"] is False