feat: enforce daily quota and surface structured 429 with settings link

openhands/app_server/app_conversation/app_conversation_router.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import hashlib
+import inspect
 import json
 import logging
 import os
@@ -490,6 +491,27 @@ async def start_app_conversation(
 ) -> AppConversationStartTask:
     await _validate_codex_credentials(start_request, user_context, secrets_store)
 
+    quota_user_id_result = user_context.get_user_id()
+    quota_user_id = (
+        await quota_user_id_result
+        if inspect.isawaitable(quota_user_id_result)
+        else quota_user_id_result
+    )
+    get_effective_org_id = getattr(user_context, 'get_effective_org_id', None)
+    quota_org_id_result = (
+        get_effective_org_id()
+        if quota_user_id and get_effective_org_id is not None
+        else None
+    )
+    quota_org_id = (
+        await quota_org_id_result
+        if inspect.isawaitable(quota_org_id_result)
+        else quota_org_id_result
+    )
+    quota_reserved = await _reserve_daily_conversation_quota(
+        quota_user_id, quota_org_id
+    )
+
     # Because we are processing after the request finishes, keep the db connection open
     set_db_session_keep_open(request.state, True)
     set_httpx_client_keep_open(request.state, True)
@@ -521,6 +543,8 @@ async def start_app_conversation(
         asyncio.create_task(_consume_remaining(async_iter, db_session, httpx_client))
         return result
     except Exception:
+        if quota_reserved and quota_user_id:
+            await _release_daily_conversation_quota(quota_user_id)
         await db_session.close()
         await httpx_client.aclose()
         raise
@@ -1156,8 +1180,15 @@ async def stream_app_conversation_start(
     await _validate_codex_credentials(request, user_context, secrets_store)
     quota_user_id = await user_context.get_user_id()
     get_effective_org_id = getattr(user_context, 'get_effective_org_id', None)
+    quota_org_id_result = (
+        get_effective_org_id()
+        if quota_user_id and get_effective_org_id is not None
+        else None
+    )
     quota_org_id = (
-        await get_effective_org_id() if get_effective_org_id is not None else None
+        await quota_org_id_result
+        if inspect.isawaitable(quota_org_id_result)
+        else quota_org_id_result
     )
     quota_reserved = await _reserve_daily_conversation_quota(
         quota_user_id, quota_org_id
@@ -2119,8 +2150,15 @@ async def _stream_app_conversation_start(
     """Stream a json list, item by item."""
     quota_user_id = await user_context.get_user_id()
     get_effective_org_id = getattr(user_context, 'get_effective_org_id', None)
+    quota_org_id_result = (
+        get_effective_org_id()
+        if quota_user_id and get_effective_org_id is not None
+        else None
+    )
     quota_org_id = (
-        await get_effective_org_id() if get_effective_org_id is not None else None
+        await quota_org_id_result
+        if inspect.isawaitable(quota_org_id_result)
+        else quota_org_id_result
     )
     quota_reserved = await _reserve_daily_conversation_quota(
         quota_user_id, quota_org_id

server/services/daily_conversation_quota_service.py

@@ -1,8 +1,9 @@
-"""Read-only daily conversation quota status for the authenticated user."""
+"""Read-only daily conversation quota status and enforcement for SaaS."""
 
 from __future__ import annotations
 
 import os
+from dataclasses import dataclass
 from datetime import UTC, date, datetime, timedelta
 from uuid import UUID
 
@@ -17,11 +18,6 @@
 from storage.user import User
 
 DEFAULT_ENV_VAR = 'OH_DAILY_CONVERSATION_LIMIT'
-
-# Sentinel stored in ``user.daily_conversation_limit`` /
-# ``org.daily_conversation_limit`` meaning "exempt -- no limit at all".
-# NULL cannot carry that meaning because NULL already means "inherit from
-# the next level down".
 EXEMPT_LIMIT = -1
 QUOTA_INCREASE_REQUEST_URL = 'https://u8mk1.share.hsforms.com/2lXOvoRtHRfeWEmba8CdOGw'
 
@@ -43,13 +39,13 @@ class QuotaStatus(BaseModel):
     reset_at: str
 
 
+@dataclass
 class DailyConversationQuotaService:
-    """Read-only quota status. Enforcement is added in a later stacked PR."""
+    """Quota status and enforcement. Enforcement is gated by a configured limit."""
 
-    def __init__(self, db_session: AsyncSession) -> None:
-        self.db_session = db_session
+    db_session: AsyncSession
 
-    async def get_status(self, user_id: str, org_id: UUID) -> QuotaStatus:
+    async def get_status(self, user_id: str, org_id: UUID | None = None) -> QuotaStatus:
         limit = await self.get_limit(user_id, org_id)
         used = await self._used(user_id, datetime.now(UTC).date())
         remaining = None if limit is None else max(limit - used, 0)
@@ -61,52 +57,45 @@ async def get_status(self, user_id: str, org_id: UUID) -> QuotaStatus:
             reset_at=reset_at,
         )
 
-    async def get_limit(self, user_id: str, org_id: UUID | None) -> int | None:
-        """Resolve the effective daily conversation limit, or None for unlimited.
-
-        Precedence (first non-NULL level wins):
+    async def get_limit(self, user_id: str, org_id: UUID | None = None) -> int | None:
+        """Resolve the effective daily conversation limit.
 
+        Precedence (first non-None wins):
         1. User-level override (``user.daily_conversation_limit``)
         2. Org-level override (``org.daily_conversation_limit``)
         3. Deployment default (``OH_DAILY_CONVERSATION_LIMIT`` env var)
 
-        NULL at a level means "inherit from the next level down"; it never
-        means exempt. To exempt (e.g. a paying SaaS org) store
-        ``EXEMPT_LIMIT`` (-1), which resolves to None at either level.
-
-        ``org_id`` must be the request's *effective* org -- resolved via
-        ``EFFECTIVE_ORG_ID`` so the API-key binding and ``X-Org-Id`` header
-        are honored. It is deliberately not derived from
-        ``user.current_org_id`` here: that is only the user's last-selected
-        org and would apply the wrong org's quota (or the wrong org's
-        exemption) whenever the request is scoped to a different one.
-        Without an effective org (``None``) only the user-level override and
-        the deployment default apply.
+        NULL at any level means "inherit from the next level down."
+        A NULL org override means the org is exempt (unlimited) — paying
+        SaaS orgs can have this set to NULL to bypass quota enforcement.
         """
         user = await self.db_session.scalar(
             select(User).where(User.id == UUID(user_id))
         )
         if user is not None and user.daily_conversation_limit is not None:
-            return self._resolve_sentinel(user.daily_conversation_limit)
+            return (
+                None
+                if user.daily_conversation_limit == EXEMPT_LIMIT
+                else user.daily_conversation_limit
+            )
 
-        return await self.get_default_limit(org_id)
+        return await self.get_default_limit(
+            org_id or (user.current_org_id if user else None)
+        )
 
     async def get_default_limit(self, org_id: UUID | None) -> int | None:
-        """Effective limit for ``org_id`` ignoring any user-level override.
-
-        Resolves the org-level override, then the deployment default.
-        Returns None when unlimited at this level (org exemption via
-        ``EXEMPT_LIMIT``, or no deployment default configured). Quota
-        increase requests cap against this value so self-service grants
-        cannot compound on top of previously granted increases.
-        """
-        org = await self.db_session.scalar(select(Org).where(Org.id == org_id))
-        if org is not None and org.daily_conversation_limit is not None:
-            return self._resolve_sentinel(org.daily_conversation_limit)
-
+        """Resolve the org-level limit and deployment default."""
+        if org_id is not None:
+            org = await self.db_session.scalar(select(Org).where(Org.id == org_id))
+            if org is not None and org.daily_conversation_limit is not None:
+                return (
+                    None
+                    if org.daily_conversation_limit == EXEMPT_LIMIT
+                    else org.daily_conversation_limit
+                )
         return configured_daily_limit()
 
-    async def reserve(self, user_id: str, org_id: UUID | None) -> bool:
+    async def reserve(self, user_id: str, org_id: UUID | None = None) -> bool:
         """Atomically increment today's usage, raising HTTP 429 at the limit."""
         limit = await self.get_limit(user_id, org_id)
         if limit is None:
@@ -161,20 +150,13 @@ def _limit_reached(limit: int, used: int, usage_date: date) -> HTTPException:
             status_code=status.HTTP_429_TOO_MANY_REQUESTS,
             detail={
                 'code': 'daily_conversation_limit_reached',
-                'message': (
-                    f'Daily conversation limit of {limit} reached. Request a quota increase at /settings/quota or {QUOTA_INCREASE_REQUEST_URL}'
-                ),
+                'message': f'Daily conversation limit of {limit} reached. Request a quota increase at /settings/quota or {QUOTA_INCREASE_REQUEST_URL}',
                 'limit': limit,
                 'used': used,
                 'reset_at': reset_at.isoformat(),
             },
         )
 
-    @staticmethod
-    def _resolve_sentinel(limit: int) -> int | None:
-        """Translate the stored exemption sentinel into "unlimited"."""
-        return None if limit == EXEMPT_LIMIT else limit
-
     @staticmethod
     def _next_reset_iso() -> str:
         """ISO timestamp of the next UTC midnight."""