diff --git a/enterprise/migrations/versions/151_add_org_daily_conversation_limit.py b/enterprise/migrations/versions/151_add_org_daily_conversation_limit.py new file mode 100644 index 000000000..db7243340 --- /dev/null +++ b/enterprise/migrations/versions/151_add_org_daily_conversation_limit.py @@ -0,0 +1,26 @@ +"""Add org-level daily conversation limit override.""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = '151' +down_revision: Union[str, None] = '150' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + bind = op.get_bind() + if bind.dialect.name != 'postgresql': + raise RuntimeError(f'Unsupported database dialect: {bind.dialect.name}') + + op.add_column( + 'org', + sa.Column('daily_conversation_limit', sa.Integer(), nullable=True), + ) + + +def downgrade() -> None: + op.drop_column('org', 'daily_conversation_limit') diff --git a/enterprise/saas_server.py b/enterprise/saas_server.py index 9443f147b..f2c26b335 100644 --- a/enterprise/saas_server.py +++ b/enterprise/saas_server.py @@ -53,7 +53,7 @@ ) from server.routes.org_profiles import router as org_profiles_router # noqa: E402 from server.routes.orgs import org_router # noqa: E402 -from server.routes.quota import quota_router # noqa: E402 +from server.routes.quota import quota_admin_router, quota_router # noqa: E402 from server.routes.readiness import readiness_router # noqa: E402 from server.routes.service import service_router # noqa: E402 from server.routes.super_admins import super_admin_router # noqa: E402 @@ -211,6 +211,7 @@ def is_saas(): analytics_events_router ) # Add routes for client-initiated analytics events base_app.include_router(quota_router) # Add routes for quota status +base_app.include_router(quota_admin_router) # Add admin routes for org-level quota base_app.add_middleware( diff --git a/enterprise/server/auth/authorization.py b/enterprise/server/auth/authorization.py index 7fae2d363..cd1b9722e 100644 --- a/enterprise/server/auth/authorization.py +++ b/enterprise/server/auth/authorization.py @@ -107,6 +107,12 @@ class Permission(str, Enum): # granted only to the ``superadmin`` super role. MANAGE_SUPER_ADMINS = 'manage_super_admins' + # Instance-level quota administration: set or clear an organization's + # daily conversation limit. Like MANAGE_SUPER_ADMINS this is an explicit + # instance-admin capability -- it is NOT implied by any org-scoped role, + # so an org owner cannot lift their own org's quota. + MANAGE_ORG_QUOTA = 'manage_org_quota' + class RoleName(str, Enum): """Role names used in the system. @@ -233,7 +239,8 @@ def super_role_name(role_name: str) -> str: SUPER_ROLE_PERMISSIONS: dict[RoleName, frozenset[Permission]] = { # Only superadmin is functional for now. It can create organizations, # provision users into a selected organization without becoming an org - # member itself, and grant/revoke the super-admin role on other users. + # member itself, grant/revoke the super-admin role on other users, and + # set org-level conversation quotas. # Additional instance-admin capabilities should be added here explicitly # as the corresponding routes are wired to permission checks. RoleName.OWNER: frozenset(), @@ -242,6 +249,7 @@ def super_role_name(role_name: str) -> str: Permission.CREATE_ORGANIZATION, Permission.PROVISION_USER, Permission.MANAGE_SUPER_ADMINS, + Permission.MANAGE_ORG_QUOTA, ] ), RoleName.MEMBER: frozenset(), diff --git a/enterprise/server/routes/quota.py b/enterprise/server/routes/quota.py index 9d431705d..8c966dc55 100644 --- a/enterprise/server/routes/quota.py +++ b/enterprise/server/routes/quota.py @@ -1,14 +1,22 @@ -"""Quota status API for the settings page.""" +"""Quota status and org-level quota management API.""" -from fastapi import APIRouter, Depends -from pydantic import BaseModel +from uuid import UUID + +from fastapi import APIRouter, Depends, HTTPException, status +from pydantic import BaseModel, Field, field_validator +from server.auth.authorization import Permission, require_permission +from server.auth.org_context import EFFECTIVE_ORG_ID, REJECT_X_ORG_ID_PATH_MISMATCH from server.services.daily_conversation_quota_service import ( + EXEMPT_LIMIT, DailyConversationQuotaService, ) +from sqlalchemy import select from storage.database import a_session_maker +from storage.org import Org from openhands.app_server.user_auth import get_user_id from openhands.app_server.utils.dependencies import get_dependencies +from openhands.app_server.utils.logger import openhands_logger as logger quota_router = APIRouter( prefix='/api/quota', tags=['Quota'], dependencies=get_dependencies() @@ -23,14 +31,109 @@ class QuotaStatusResponse(BaseModel): @quota_router.get('/status', response_model=QuotaStatusResponse) -async def get_quota_status(user_id: str = Depends(get_user_id)) -> QuotaStatusResponse: - """Return the authenticated user's daily conversation quota status.""" +async def get_quota_status( + user_id: str = Depends(get_user_id), + effective_org_id: UUID = EFFECTIVE_ORG_ID, +) -> QuotaStatusResponse: + """Return the authenticated user's daily conversation quota status. + + Scoped to the request's effective org (``X-Org-Id`` / API-key binding), + so a user in several orgs sees the quota of the org they are actually + working in. + """ async with a_session_maker() as session: service = DailyConversationQuotaService(session) - status = await service.get_status(user_id) + quota = await service.get_status(user_id, effective_org_id) return QuotaStatusResponse( - daily_limit=status.daily_limit, - used_today=status.used_today, - remaining=status.remaining, - reset_at=status.reset_at, + daily_limit=quota.daily_limit, + used_today=quota.used_today, + remaining=quota.remaining, + reset_at=quota.reset_at, + ) + + +# --- Org-level quota management (admin) --- + +quota_admin_router = APIRouter( + prefix='/api/admin/quota', + tags=['Admin'], + dependencies=[*get_dependencies(), REJECT_X_ORG_ID_PATH_MISMATCH], +) + + +class OrgQuotaUpdateRequest(BaseModel): + daily_conversation_limit: int | None = Field( + description=( + 'NULL to inherit the deployment default, -1 to exempt the org ' + 'entirely, or a positive integer for an org-specific limit.' + ), ) + + @field_validator('daily_conversation_limit') + @classmethod + def _reject_meaningless_limits(cls, value: int | None) -> int | None: + """Allow only NULL, the exemption sentinel, and positive limits. + + 0 and values below -1 are rejected rather than stored: they are not + meaningful quotas, but they resolve to a limit the org can never + satisfy, silently blocking every member with no error anywhere. A + mistyped '-11' for '-1' would otherwise do the exact opposite of the + intended exemption. + """ + if value is None or value == EXEMPT_LIMIT or value > 0: + return value + raise ValueError( + f'daily_conversation_limit must be null (inherit), {EXEMPT_LIMIT} ' + f'(exempt), or a positive integer; got {value}' + ) + + +class OrgQuotaResponse(BaseModel): + org_id: str + org_name: str + daily_conversation_limit: int | None + + +@quota_admin_router.put('/orgs/{org_id}/quota', response_model=OrgQuotaResponse) +async def set_org_quota( + org_id: UUID, + body: OrgQuotaUpdateRequest, + caller_user_id: str = Depends(require_permission(Permission.MANAGE_ORG_QUOTA)), +) -> OrgQuotaResponse: + """Set or clear an org-level daily conversation limit override. + + Requires the instance-level ``MANAGE_ORG_QUOTA`` permission, which is + granted only to the superadmin super role. Going through + ``require_permission`` (rather than checking the super role inline) also + enforces the API-key organization binding, so a key bound to one org + cannot edit another org's quota. + + Set to -1 to exempt the org entirely (unlimited). + Set to NULL to inherit the deployment default. + Set to a positive integer for an org-specific limit. + """ + async with a_session_maker() as session: + org = await session.scalar(select(Org).where(Org.id == org_id)) + if org is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail='Organization not found', + ) + org.daily_conversation_limit = body.daily_conversation_limit + await session.commit() + + logger.info( + 'org_quota:set', + extra={ + 'caller_user_id': caller_user_id, + 'org_id': str(org.id), + 'daily_conversation_limit': org.daily_conversation_limit, + 'exempt': org.daily_conversation_limit == EXEMPT_LIMIT, + }, + ) + + return OrgQuotaResponse( + org_id=str(org.id), + org_name=org.name, + daily_conversation_limit=org.daily_conversation_limit, + ) diff --git a/enterprise/server/services/daily_conversation_quota_service.py b/enterprise/server/services/daily_conversation_quota_service.py index 9f7cc16c4..722ed4032 100644 --- a/enterprise/server/services/daily_conversation_quota_service.py +++ b/enterprise/server/services/daily_conversation_quota_service.py @@ -10,10 +10,17 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from storage.daily_conversation_usage import DailyConversationUsage +from storage.org import Org 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 + def configured_daily_limit() -> int | None: """Read the deployment default; unset and blank mean unlimited.""" @@ -38,8 +45,8 @@ class DailyConversationQuotaService: def __init__(self, db_session: AsyncSession) -> None: self.db_session = db_session - async def get_status(self, user_id: str) -> QuotaStatus: - limit = await self.get_limit(user_id) + async def get_status(self, user_id: str, org_id: UUID) -> 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) reset_at = self._next_reset_iso() @@ -50,13 +57,42 @@ async def get_status(self, user_id: str) -> QuotaStatus: reset_at=reset_at, ) - async def get_limit(self, user_id: str) -> int | None: + async def get_limit(self, user_id: str, org_id: UUID) -> int | None: + """Resolve the effective daily conversation limit, or None for unlimited. + + Precedence (first non-NULL level 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. + """ user = await self.db_session.scalar( select(User).where(User.id == UUID(user_id)) ) - if user is None or user.daily_conversation_limit is None: - return configured_daily_limit() - return user.daily_conversation_limit + if user is not None and user.daily_conversation_limit is not None: + return self._resolve_sentinel(user.daily_conversation_limit) + + 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) + + return configured_daily_limit() + + @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: diff --git a/enterprise/storage/org.py b/enterprise/storage/org.py index ed04e5b2e..400648730 100644 --- a/enterprise/storage/org.py +++ b/enterprise/storage/org.py @@ -69,6 +69,10 @@ class Org(Base): # Set by completed billing sessions or when positive org credits are detected. byor_export_enabled: Mapped[bool] = mapped_column(nullable=False, default=False) sandbox_grouping_strategy: Mapped[str | None] = mapped_column(String, nullable=True) + # NULL inherits the deployment-wide daily conversation limit. + # Set to -1 to exempt this org from daily conversation limits entirely + # (for paying SaaS customers). Any other integer is the org-specific limit. + daily_conversation_limit: Mapped[int | None] = mapped_column(nullable=True) # Encrypted column for LLM profiles (contains API keys) llm_profiles: Mapped[dict[str, Any] | None] = mapped_column( EncryptedJSON, nullable=True diff --git a/enterprise/tests/unit/server/routes/test_quota_admin.py b/enterprise/tests/unit/server/routes/test_quota_admin.py new file mode 100644 index 000000000..41785864c --- /dev/null +++ b/enterprise/tests/unit/server/routes/test_quota_admin.py @@ -0,0 +1,274 @@ +"""Route-level tests for the org quota administration API. + +These exercise the FastAPI route through ``require_permission`` and stub +the DB session at the route module boundary. Authorization is faked the +same way ``test_super_admins.py`` does: short-circuit the org-role lookup +to ``None`` and stack a ``superadmin`` super role on top of the +conftest-level ``get_user_super_role -> None`` default. +""" + +import uuid +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient +from server.routes.quota import quota_admin_router + +from openhands.app_server.user_auth import get_user_id + +CALLER_USER_ID = str(uuid.uuid4()) +ORG_ID = uuid.uuid4() + + +@pytest.fixture +def mock_app(): + app = FastAPI() + app.include_router(quota_admin_router) + app.dependency_overrides[get_user_id] = lambda: CALLER_USER_ID + return app + + +@pytest.fixture +def grant_manage_org_quota(): + """Make ``MANAGE_ORG_QUOTA`` succeed by faking a ``superadmin`` role.""" + superadmin = MagicMock() + superadmin.name = 'admin' + with ( + patch( + 'server.auth.authorization.get_user_org_role', + AsyncMock(return_value=None), + ), + patch( + 'server.auth.authorization.get_user_super_role', + AsyncMock(return_value=superadmin), + ), + ): + yield + + +@pytest.fixture +def deny_manage_org_quota(): + """Caller is an org owner but holds no super role.""" + owner = MagicMock() + owner.name = 'owner' + with patch( + 'server.auth.authorization.get_user_org_role', + AsyncMock(return_value=owner), + ): + yield + + +def _client(app): + return AsyncClient(transport=ASGITransport(app=app), base_url='http://test') + + +def _fake_org(limit: int | None = None): + org = MagicMock() + org.id = ORG_ID + org.name = 'Acme' + org.daily_conversation_limit = limit + return org + + +@asynccontextmanager +async def _session_yielding(org): + session = AsyncMock() + session.scalar = AsyncMock(return_value=org) + yield session + + +def _patch_session(org): + return patch( + 'server.routes.quota.a_session_maker', + lambda **kwargs: _session_yielding(org), + ) + + +def _url(org_id=ORG_ID): + return f'/api/admin/quota/orgs/{org_id}/quota' + + +@pytest.mark.asyncio +async def test_superadmin_can_set_org_limit(mock_app, grant_manage_org_quota): + org = _fake_org() + with _patch_session(org): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': 25}) + + assert resp.status_code == 200 + assert resp.json() == { + 'org_id': str(ORG_ID), + 'org_name': 'Acme', + 'daily_conversation_limit': 25, + } + assert org.daily_conversation_limit == 25 + + +@pytest.mark.asyncio +async def test_superadmin_can_exempt_org(mock_app, grant_manage_org_quota): + org = _fake_org(limit=10) + with _patch_session(org): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': -1}) + + assert resp.status_code == 200 + assert resp.json()['daily_conversation_limit'] == -1 + + +@pytest.mark.asyncio +async def test_superadmin_can_clear_override(mock_app, grant_manage_org_quota): + org = _fake_org(limit=10) + with _patch_session(org): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': None}) + + assert resp.status_code == 200 + assert resp.json()['daily_conversation_limit'] is None + + +@pytest.mark.asyncio +async def test_unknown_org_returns_404(mock_app, grant_manage_org_quota): + with _patch_session(None): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': 5}) + + assert resp.status_code == 404 + assert resp.json()['detail'] == 'Organization not found' + + +@pytest.mark.asyncio +async def test_org_owner_without_super_role_is_denied(mock_app, deny_manage_org_quota): + """MANAGE_ORG_QUOTA is instance-level: no org-scoped role grants it.""" + org = _fake_org() + with _patch_session(org): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': 5}) + + assert resp.status_code == 403 + assert org.daily_conversation_limit is None + + +@pytest.mark.asyncio +async def test_org_bound_api_key_cannot_edit_another_org( + mock_app, grant_manage_org_quota +): + """An API key bound to one org must not reach another org's quota. + + Regression guard for the hand-rolled super-role check this route used + to carry, which skipped ``require_permission``'s API-key organization + binding entirely. + """ + other_org = uuid.uuid4() + org = _fake_org() + with ( + patch( + 'server.auth.authorization.get_api_key_org_id_from_request', + AsyncMock(return_value=other_org), + ), + _patch_session(org), + ): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': 5}) + + assert resp.status_code == 403 + assert 'not authorized for this organization' in resp.json()['detail'] + assert org.daily_conversation_limit is None + + +@pytest.mark.asyncio +async def test_org_bound_api_key_can_edit_its_own_org(mock_app, grant_manage_org_quota): + org = _fake_org() + with ( + patch( + 'server.auth.authorization.get_api_key_org_id_from_request', + AsyncMock(return_value=ORG_ID), + ), + _patch_session(org), + ): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': 5}) + + assert resp.status_code == 200 + assert org.daily_conversation_limit == 5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize('limit', [0, -2, -11, -100]) +async def test_meaningless_limits_are_rejected(mock_app, grant_manage_org_quota, limit): + """0 and values below -1 would silently hard-block the whole org.""" + org = _fake_org() + with _patch_session(org): + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': limit}) + + assert resp.status_code == 422 + assert org.daily_conversation_limit is None + + +@pytest.mark.asyncio +async def test_malformed_org_id_is_rejected(mock_app, grant_manage_org_quota): + """A non-UUID org id is a 422, not an unhandled 500.""" + async with _client(mock_app) as client: + resp = await client.put( + _url('not-a-uuid'), json={'daily_conversation_limit': 5} + ) + + assert resp.status_code == 422 + + +@pytest.mark.asyncio +async def test_conflicting_x_org_id_header_is_rejected( + mock_app, grant_manage_org_quota +): + """The path pins the org; a contradictory X-Org-Id header is a 400.""" + org = _fake_org() + with _patch_session(org): + async with _client(mock_app) as client: + resp = await client.put( + _url(), + json={'daily_conversation_limit': 5}, + headers={'X-Org-Id': str(uuid.uuid4())}, + ) + + assert resp.status_code == 400 + assert org.daily_conversation_limit is None + + +@pytest.mark.asyncio +async def test_successful_change_is_logged_with_the_caller( + mock_app, grant_manage_org_quota +): + """The mutation is attributable: who changed which org's quota, to what. + + Exempting an org is revenue-affecting, so an unattributable write is + not good enough -- this mirrors ``super_admins:grant`` / ``:revoke``. + """ + org = _fake_org() + with _patch_session(org), patch('server.routes.quota.logger') as log: + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': -1}) + + assert resp.status_code == 200 + log.info.assert_called_once() + event, kwargs = log.info.call_args[0][0], log.info.call_args[1] + assert event == 'org_quota:set' + assert kwargs['extra'] == { + 'caller_user_id': CALLER_USER_ID, + 'org_id': str(ORG_ID), + 'daily_conversation_limit': -1, + 'exempt': True, + } + + +@pytest.mark.asyncio +async def test_denied_request_logs_no_change(mock_app, deny_manage_org_quota): + """A rejected call must not leave a record implying a change happened.""" + org = _fake_org() + with _patch_session(org), patch('server.routes.quota.logger') as log: + async with _client(mock_app) as client: + resp = await client.put(_url(), json={'daily_conversation_limit': 5}) + + assert resp.status_code == 403 + log.info.assert_not_called() diff --git a/enterprise/tests/unit/server/routes/test_quota_status.py b/enterprise/tests/unit/server/routes/test_quota_status.py new file mode 100644 index 000000000..ee252ea0f --- /dev/null +++ b/enterprise/tests/unit/server/routes/test_quota_status.py @@ -0,0 +1,98 @@ +"""Route-level tests for the quota status endpoint. + +The point of interest is org scoping: the route must hand the service the +request's *effective* org (``X-Org-Id`` / API-key binding, resolved by +``EFFECTIVE_ORG_ID``) rather than letting the service fall back to the +user's last-selected ``current_org_id``. +""" + +import uuid +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient +from server.auth.org_context import resolve_effective_org_id +from server.routes.quota import quota_router +from server.services.daily_conversation_quota_service import QuotaStatus + +from openhands.app_server.user_auth import get_user_id + +CALLER_USER_ID = str(uuid.uuid4()) +EFFECTIVE_ORG = uuid.uuid4() + + +@asynccontextmanager +async def _session(): + yield AsyncMock() + + +@pytest.fixture +def mock_app(): + app = FastAPI() + app.include_router(quota_router) + app.dependency_overrides[get_user_id] = lambda: CALLER_USER_ID + app.dependency_overrides[resolve_effective_org_id] = lambda: EFFECTIVE_ORG + return app + + +def _client(app): + return AsyncClient(transport=ASGITransport(app=app), base_url='http://test') + + +@pytest.mark.asyncio +async def test_status_is_scoped_to_the_effective_org(mock_app): + status = QuotaStatus( + daily_limit=15, + used_today=3, + remaining=12, + reset_at='2026-08-26T00:00:00+00:00', + ) + get_status = AsyncMock(return_value=status) + + with ( + patch('server.routes.quota.a_session_maker', lambda **kwargs: _session()), + patch( + 'server.routes.quota.DailyConversationQuotaService.get_status', + get_status, + ), + ): + async with _client(mock_app) as client: + resp = await client.get('/api/quota/status') + + assert resp.status_code == 200 + assert resp.json() == { + 'daily_limit': 15, + 'used_today': 3, + 'remaining': 12, + 'reset_at': '2026-08-26T00:00:00+00:00', + } + # The effective org -- not the user's current_org_id -- reaches the service. + get_status.assert_awaited_once_with(CALLER_USER_ID, EFFECTIVE_ORG) + + +@pytest.mark.asyncio +async def test_status_reports_unlimited_without_a_limit(mock_app): + status = QuotaStatus( + daily_limit=None, + used_today=7, + remaining=None, + reset_at='2026-08-26T00:00:00+00:00', + ) + + with ( + patch('server.routes.quota.a_session_maker', lambda **kwargs: _session()), + patch( + 'server.routes.quota.DailyConversationQuotaService.get_status', + AsyncMock(return_value=status), + ), + ): + async with _client(mock_app) as client: + resp = await client.get('/api/quota/status') + + assert resp.status_code == 200 + body = resp.json() + assert body['daily_limit'] is None + assert body['remaining'] is None + assert body['used_today'] == 7 diff --git a/enterprise/tests/unit/test_authorization.py b/enterprise/tests/unit/test_authorization.py index 1c7d5cfd2..27c7d2e9c 100644 --- a/enterprise/tests/unit/test_authorization.py +++ b/enterprise/tests/unit/test_authorization.py @@ -1473,6 +1473,7 @@ def test_super_role_permissions_are_explicit(self): Permission.CREATE_ORGANIZATION, Permission.PROVISION_USER, Permission.MANAGE_SUPER_ADMINS, + Permission.MANAGE_ORG_QUOTA, ] ) assert SUPER_ROLE_PERMISSIONS[RoleName.MEMBER] == frozenset() @@ -1520,6 +1521,23 @@ def test_only_superadmin_grants_instance_org_management(self): assert Permission.PROVISION_USER in SUPER_ROLE_PERMISSIONS[RoleName.ADMIN] assert Permission.PROVISION_USER not in SUPER_ROLE_PERMISSIONS[RoleName.MEMBER] + def test_manage_org_quota_is_superadmin_only(self): + """ + GIVEN: SUPER_ROLE_PERMISSIONS and ROLE_PERMISSIONS + WHEN: looking up MANAGE_ORG_QUOTA + THEN: only the superadmin super role carries it -- no org-scoped + role does, so an org owner cannot lift their own org's quota. + """ + assert Permission.MANAGE_ORG_QUOTA in SUPER_ROLE_PERMISSIONS[RoleName.ADMIN] + assert Permission.MANAGE_ORG_QUOTA not in SUPER_ROLE_PERMISSIONS[RoleName.OWNER] + assert ( + Permission.MANAGE_ORG_QUOTA not in SUPER_ROLE_PERMISSIONS[RoleName.MEMBER] + ) + for role_name, permissions in ROLE_PERMISSIONS.items(): + assert Permission.MANAGE_ORG_QUOTA not in permissions, ( + f'org-scoped role {role_name} must not grant MANAGE_ORG_QUOTA' + ) + class TestHasPermissionSuper: """Tests for ``has_permission`` with the ``is_super`` flag.""" diff --git a/enterprise/tests/unit/test_daily_conversation_quota_service.py b/enterprise/tests/unit/test_daily_conversation_quota_service.py index 598c08843..32348f071 100644 --- a/enterprise/tests/unit/test_daily_conversation_quota_service.py +++ b/enterprise/tests/unit/test_daily_conversation_quota_service.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch from uuid import uuid4 import pytest @@ -9,6 +9,7 @@ ) USER_ID = str(uuid4()) +ORG_ID = uuid4() @pytest.mark.asyncio @@ -16,12 +17,14 @@ async def test_get_status_unlimited(): """When limit is None (unlimited), remaining is None and reset_at is next UTC midnight.""" session = AsyncMock() session.scalar.side_effect = [ - SimpleNamespace(daily_conversation_limit=None), + SimpleNamespace(daily_conversation_limit=None), # user: no override + SimpleNamespace(daily_conversation_limit=None), # org: no override None, # no usage record for today ] - service = DailyConversationQuotaService(session) - result = await service.get_status(USER_ID) + with patch.dict('os.environ', {}, clear=True): + service = DailyConversationQuotaService(session) + result = await service.get_status(USER_ID, ORG_ID) assert isinstance(result, QuotaStatus) assert result.daily_limit is None @@ -42,7 +45,7 @@ async def test_get_status_with_limit_and_usage(): ] service = DailyConversationQuotaService(session) - result = await service.get_status(USER_ID) + result = await service.get_status(USER_ID, ORG_ID) assert result.daily_limit == 20 assert result.used_today == 5 @@ -60,7 +63,7 @@ async def test_get_status_remaining_floor_zero(): ] service = DailyConversationQuotaService(session) - result = await service.get_status(USER_ID) + result = await service.get_status(USER_ID, ORG_ID) assert result.daily_limit == 10 assert result.used_today == 12 @@ -77,8 +80,50 @@ async def test_get_status_no_usage_today(): ] service = DailyConversationQuotaService(session) - result = await service.get_status(USER_ID) + result = await service.get_status(USER_ID, ORG_ID) assert result.daily_limit == 20 assert result.used_today == 0 assert result.remaining == 20 + + +@pytest.mark.asyncio +async def test_get_status_resolves_org_limit_through_full_query_sequence(): + """The production path: user override NULL -> org override -> usage. + + Pins the real three-query sequence (user, org, usage) that every + request takes, so an ordering regression in ``get_status`` is caught. + """ + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None), # user: inherit + SimpleNamespace(daily_conversation_limit=30), # org: org-specific limit + SimpleNamespace(conversation_count=4), # usage today + ] + + service = DailyConversationQuotaService(session) + result = await service.get_status(USER_ID, ORG_ID) + + assert session.scalar.await_count == 3 + assert result.daily_limit == 30 + assert result.used_today == 4 + assert result.remaining == 26 + + +@pytest.mark.asyncio +async def test_get_status_org_exemption_reports_unlimited(): + """An exempt org (-1) surfaces as unlimited rather than a negative limit.""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None), # user: inherit + SimpleNamespace(daily_conversation_limit=-1), # org: exempt + SimpleNamespace(conversation_count=99), # usage today + ] + + with patch.dict('os.environ', {'OH_DAILY_CONVERSATION_LIMIT': '20'}): + service = DailyConversationQuotaService(session) + result = await service.get_status(USER_ID, ORG_ID) + + assert result.daily_limit is None + assert result.used_today == 99 + assert result.remaining is None diff --git a/enterprise/tests/unit/test_migration_graph.py b/enterprise/tests/unit/test_migration_graph.py new file mode 100644 index 000000000..5536067c2 --- /dev/null +++ b/enterprise/tests/unit/test_migration_graph.py @@ -0,0 +1,91 @@ +"""Integrity checks for the Alembic revision graph. + +Migrations here are numbered sequentially, so two branches opened off the +same parent independently pick the *same* next number. Each branch is +valid on its own -- and passes CI on its own -- but the second one to +merge lands a duplicate revision id and a second head on ``main``, which +breaks ``alembic upgrade head`` for every deploy. + +These tests make that collision fail loudly on the second branch instead +of silently on ``main`` after the merge. +""" + +import re +from pathlib import Path + +VERSIONS_DIR = Path(__file__).resolve().parents[2] / 'migrations' / 'versions' + +_REVISION_RE = re.compile(r'^revision(?::\s*[^=]+)?\s*=\s*[\'"]([^\'"]+)[\'"]', re.M) +_DOWN_REVISION_RE = re.compile( + r'^down_revision(?::\s*[^=]+)?\s*=\s*(?:[\'"]([^\'"]+)[\'"]|None)', re.M +) + + +def _migrations() -> dict[str, tuple[str, str | None]]: + """Map each migration file to its ``(revision, down_revision)`` pair.""" + found: dict[str, tuple[str, str | None]] = {} + for path in sorted(VERSIONS_DIR.glob('*.py')): + if path.name == '__init__.py': + continue + source = path.read_text() + revision_match = _REVISION_RE.search(source) + down_match = _DOWN_REVISION_RE.search(source) + assert revision_match is not None, f'{path.name} declares no revision' + assert down_match is not None, f'{path.name} declares no down_revision' + found[path.name] = (revision_match.group(1), down_match.group(1)) + return found + + +def test_revision_ids_are_unique(): + """Two migrations must never claim the same revision id.""" + by_revision: dict[str, list[str]] = {} + for filename, (revision, _) in _migrations().items(): + by_revision.setdefault(revision, []).append(filename) + + duplicates = {rev: files for rev, files in by_revision.items() if len(files) > 1} + assert not duplicates, ( + 'Duplicate Alembic revision ids: ' + f'{duplicates}. Another branch already merged this number -- ' + 'renumber this migration to the next free revision.' + ) + + +def test_revision_graph_has_a_single_head(): + """Exactly one migration must be unreferenced as a parent.""" + migrations = _migrations() + revisions = {revision for revision, _ in migrations.values()} + parents = {down for _, down in migrations.values() if down is not None} + + heads = sorted(revisions - parents) + assert len(heads) == 1, ( + f'Expected exactly one Alembic head, found {len(heads)}: {heads}. ' + 'Concurrent migrations branched off the same parent -- rebase one ' + 'onto the other so the chain stays linear.' + ) + + +def test_every_down_revision_exists(): + """No migration may point at a parent that is not in the tree.""" + migrations = _migrations() + revisions = {revision for revision, _ in migrations.values()} + + orphans = { + filename: down + for filename, (_, down) in migrations.items() + if down is not None and down not in revisions + } + assert not orphans, f'Migrations reference missing parents: {orphans}' + + +def test_no_two_migrations_share_a_parent(): + """A shared parent is a fork in the chain, even if one head still wins.""" + by_parent: dict[str, list[str]] = {} + for filename, (_, down) in _migrations().items(): + if down is not None: + by_parent.setdefault(down, []).append(filename) + + forks = {down: files for down, files in by_parent.items() if len(files) > 1} + assert not forks, ( + f'Multiple migrations share a down_revision: {forks}. ' + 'Chain them sequentially instead of branching.' + ) diff --git a/enterprise/tests/unit/test_org_level_quota.py b/enterprise/tests/unit/test_org_level_quota.py new file mode 100644 index 000000000..7bce146d9 --- /dev/null +++ b/enterprise/tests/unit/test_org_level_quota.py @@ -0,0 +1,140 @@ +"""Tests for org-level daily conversation quota resolution.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +import pytest +from server.services.daily_conversation_quota_service import ( + DailyConversationQuotaService, +) + +USER_ID = str(uuid4()) +ORG_ID = uuid4() + + +@pytest.mark.asyncio +async def test_user_override_takes_precedence(): + """User-level override wins over org and env.""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=50), # user + SimpleNamespace(daily_conversation_limit=100), # org (never reached) + ] + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) == 50 + # The org row is not even queried once the user override resolves. + assert session.scalar.await_count == 1 + + +@pytest.mark.asyncio +async def test_org_override_when_user_is_null(): + """When user override is NULL, org override applies.""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None), # user + SimpleNamespace(daily_conversation_limit=100), # org + ] + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) == 100 + + +@pytest.mark.asyncio +async def test_org_exempt_means_unlimited(): + """When org limit is -1, the user is exempt (returns None = unlimited).""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None), # user + SimpleNamespace(daily_conversation_limit=-1), # exempt + ] + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) is None + + +@pytest.mark.asyncio +async def test_user_exempt_means_unlimited(): + """-1 means exempt at the user level too, not a literal limit of -1.""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=-1), # user: exempt + SimpleNamespace(daily_conversation_limit=5), # org (never reached) + ] + with patch.dict('os.environ', {'OH_DAILY_CONVERSATION_LIMIT': '20'}): + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) is None + + +@pytest.mark.asyncio +async def test_falls_back_to_env_when_both_null(): + """When both user and org overrides are NULL, env default applies.""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None), # user + SimpleNamespace(daily_conversation_limit=None), # org: not set + ] + with patch.dict('os.environ', {'OH_DAILY_CONVERSATION_LIMIT': '20'}): + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) == 20 + + +@pytest.mark.asyncio +async def test_falls_back_to_none_when_unset(): + """When everything is unset, returns None (unlimited).""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None), # user + SimpleNamespace(daily_conversation_limit=None), # org + ] + with patch.dict('os.environ', {}, clear=True): + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) is None + + +@pytest.mark.asyncio +async def test_user_null_org_exempt_takes_precedence_over_env(): + """Org exemption (-1) takes precedence over env default.""" + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None), # user + SimpleNamespace(daily_conversation_limit=-1), # exempt + ] + with patch.dict('os.environ', {'OH_DAILY_CONVERSATION_LIMIT': '20'}): + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) is None + + +@pytest.mark.asyncio +async def test_missing_user_row_still_resolves_org_override(): + """A missing user row falls through to the org, not straight to env.""" + session = AsyncMock() + session.scalar.side_effect = [ + None, # user row absent + SimpleNamespace(daily_conversation_limit=7), # org + ] + with patch.dict('os.environ', {'OH_DAILY_CONVERSATION_LIMIT': '20'}): + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) == 7 + + +@pytest.mark.asyncio +async def test_limit_is_resolved_for_the_org_passed_in(): + """The org queried is the caller-supplied effective org. + + Regression guard: resolution must not fall back to the user's + ``current_org_id``, which is only their last-selected org and would + apply the wrong org's quota whenever the request is scoped elsewhere + via ``X-Org-Id`` or an org-bound API key. + """ + other_org = uuid4() + session = AsyncMock() + session.scalar.side_effect = [ + SimpleNamespace(daily_conversation_limit=None, current_org_id=other_org), + SimpleNamespace(daily_conversation_limit=42), + ] + service = DailyConversationQuotaService(session) + assert await service.get_limit(USER_ID, ORG_ID) == 42 + + org_query = session.scalar.await_args_list[1].args[0] + compiled = str(org_query.compile(compile_kwargs={'literal_binds': True})) + assert ORG_ID.hex in compiled + assert other_org.hex not in compiled