diff --git a/openhands/app_server/app_conversation/app_conversation_router.py b/openhands/app_server/app_conversation/app_conversation_router.py index 06b0f2012..f8289a707 100644 --- a/openhands/app_server/app_conversation/app_conversation_router.py +++ b/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 diff --git a/server/services/daily_conversation_quota_service.py b/server/services/daily_conversation_quota_service.py index e09ecaf5a..afeea8707 100644 --- a/server/services/daily_conversation_quota_service.py +++ b/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."""