fix(usage): hold service accounts to global budgets only (#15030)

backend/ee/onyx/server/query_and_chat/token_limit.py

@@ -6,6 +6,7 @@ from sqlalchemy.orm import Session
 from onyx.configs.constants import TokenRateLimitScope
 from onyx.db.api_key import is_api_key_email_address
 from onyx.db.engine.sql_engine import get_session_with_current_tenant
+from onyx.db.enums import AccountType
 from onyx.db.models import User
 from onyx.db.token_limit import (
     fetch_all_user_token_rate_limits,
@@ -37,8 +38,9 @@ def _check_token_rate_limits(user: User) -> None:
     if user.is_anonymous:
         _user_is_rate_limited_by_global()
 
-    elif is_api_key_email_address(user.email):
-        # API keys are only rate limited by global settings
+    elif user.account_type == AccountType.SERVICE_ACCOUNT or is_api_key_email_address(
+        user.email
+    ):
         _user_is_rate_limited_by_global()
 
     else:

backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py

@@ -346,6 +346,16 @@ def handle_regular_answer(
                 send_as_ephemeral=send_as_ephemeral,
             )
         finally:
+            if feedback_reminder_id and message_info.sender_id:
+                try:
+                    client.chat_deleteScheduledMessage(
+                        channel=message_info.sender_id,
+                        scheduled_message_id=feedback_reminder_id,
+                    )
+                except Exception:
+                    logger.exception(
+                        "Unable to delete scheduled feedback reminder after Slack budget reply"
+                    )
             if not is_slash_command:
                 update_emote_react(
                     emoji=ONYX_BOT_REACT_EMOJI,
@@ -354,7 +364,7 @@ def handle_regular_answer(
                     remove=True,
                     client=client,
                 )
-        return True
+        return False
 
     try:
         filters = BaseFilters(

backend/tests/external_dependency_unit/ee/onyx/server/query_and_chat/test_service_account_budget_scope.py

@@ -0,0 +1,121 @@
+from collections.abc import Generator
+from datetime import datetime, timezone
+
+import pytest
+from sqlalchemy import delete
+from sqlalchemy.orm import Session
+
+import ee.onyx.server.query_and_chat.token_limit as ee_token_limit
+import onyx.server.query_and_chat.token_limit as token_limit
+from onyx.configs.constants import TokenRateLimitScope
+from onyx.db.enums import AccountType
+from onyx.db.llm_usage import LLMUsageRecord
+from onyx.db.models import TokenRateLimit, User, UserUsage
+from onyx.db.user_usage import get_token_window_start, record_user_usage
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
+from onyx.tracing.flows import LLMFlow
+from tests.external_dependency_unit.conftest import create_test_user, delete_test_user
+
+pytestmark = pytest.mark.usefixtures("tenant_context")
+
+
+def _token_limit(scope: TokenRateLimitScope) -> TokenRateLimit:
+    return TokenRateLimit(
+        enabled=True,
+        token_budget=1,
+        cost_budget_cents=None,
+        period_hours=24,
+        scope=scope,
+    )
+
+
+def _record_over_budget_usage(db_session: Session, user: User) -> None:
+    record_user_usage(
+        db_session=db_session,
+        user_id=str(user.id),
+        usage=LLMUsageRecord(
+            model="service-account-budget-test-model",
+            flow=LLMFlow.CHAT_RESPONSE.value,
+            provider=None,
+            input_tokens=1_500,
+            output_tokens=500,
+            cache_read_tokens=0,
+            cost_cents=0.0,
+            window_start=get_token_window_start(datetime.now(timezone.utc), 24),
+        ),
+    )
+    db_session.commit()
+
+
+@pytest.fixture
+def users_over_a_per_user_budget(
+    db_session: Session, monkeypatch: pytest.MonkeyPatch
+) -> Generator[tuple[User, User], None, None]:
+    standard = create_test_user(db_session, "budget-scope-standard")
+    service = create_test_user(
+        db_session, "budget-scope-service", account_type=AccountType.SERVICE_ACCOUNT
+    )
+    for user in (standard, service):
+        _record_over_budget_usage(db_session, user)
+
+    monkeypatch.setattr(
+        ee_token_limit,
+        "fetch_all_user_token_rate_limits",
+        lambda **_: [_token_limit(TokenRateLimitScope.USER)],
+    )
+    monkeypatch.setattr(
+        ee_token_limit, "_user_is_rate_limited_by_group", lambda _: None
+    )
+    monkeypatch.setattr(
+        token_limit, "fetch_all_global_token_rate_limits", lambda **_: []
+    )
+
+    yield standard, service
+
+    db_session.rollback()
+    db_session.execute(
+        delete(UserUsage).where(UserUsage.user_id.in_([standard.id, service.id]))
+    )
+    delete_test_user(db_session, standard, service)
+    db_session.commit()
+
+
+def test_standard_user_over_per_user_budget_is_limited(
+    users_over_a_per_user_budget: tuple[User, User],
+) -> None:
+    standard, _ = users_over_a_per_user_budget
+
+    with pytest.raises(OnyxError) as exc_info:
+        ee_token_limit._check_token_rate_limits(standard)
+
+    assert exc_info.value.error_code is OnyxErrorCode.RATE_LIMITED
+    assert exc_info.value.extra is not None
+    assert exc_info.value.extra["scope"] == TokenRateLimitScope.USER.value
+
+
+def test_service_account_is_not_held_to_per_user_budget(
+    users_over_a_per_user_budget: tuple[User, User],
+) -> None:
+    _, service = users_over_a_per_user_budget
+
+    ee_token_limit._check_token_rate_limits(service)
+
+
+def test_service_account_is_still_held_to_global_budget(
+    users_over_a_per_user_budget: tuple[User, User],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _, service = users_over_a_per_user_budget
+    monkeypatch.setattr(
+        token_limit,
+        "fetch_all_global_token_rate_limits",
+        lambda **_: [_token_limit(TokenRateLimitScope.GLOBAL)],
+    )
+
+    with pytest.raises(OnyxError) as exc_info:
+        ee_token_limit._check_token_rate_limits(service)
+
+    assert exc_info.value.error_code is OnyxErrorCode.RATE_LIMITED
+    assert exc_info.value.extra is not None
+    assert exc_info.value.extra["scope"] == TokenRateLimitScope.GLOBAL.value

backend/tests/external_dependency_unit/server/test_llm_entrypoint_budgets.py

@@ -198,6 +198,7 @@ def test_slack_bot_replies_with_budget_message_instead_of_answering(
         ),
     )
 
+    client = MagicMock()
     with (
         patch(
             f"{_HANDLE_REGULAR_ANSWER}.get_user_by_email",
@@ -227,11 +228,11 @@ def test_slack_bot_replies_with_budget_message_instead_of_answering(
             message_info=message_info,
             slack_channel_config=slack_channel_config,
             receiver_ids=None,
-            client=MagicMock(),
+            client=client,
             channel="C123",
             logger=MagicMock(),
             db_session=db_session,
-            feedback_reminder_id=None,
+            feedback_reminder_id="scheduled-reminder",
         )
 
     mock_update_react.assert_called_once()
@@ -239,5 +240,8 @@ def test_slack_bot_replies_with_budget_message_instead_of_answering(
     mock_stream.assert_not_called()
     mock_respond.assert_called_once()
     assert mock_respond.call_args.kwargs["text"] == _BUDGET_MESSAGE
+    client.chat_deleteScheduledMessage.assert_called_once_with(
+        channel="U123", scheduled_message_id="scheduled-reminder"
+    )
     if not reply_fails:
-        assert result is True
+        assert result is False

backend/tests/unit/ee/onyx/server/gateway/test_gateway_cost_cap.py

@@ -3,21 +3,24 @@
 The chat route runs `check_llm_cost_limit_for_provider` for every provider it
 calls. The gateway's generating routes did not, so a tenant over its weekly cap
 could keep spending Onyx's managed keys through `/v1/*`. The cap itself is
-real here; only the usage lookup is stubbed to report "over the cap".
+real here.
 """
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Generator
 from typing import Any, cast
 from unittest.mock import MagicMock, patch
 
 import pytest
-from fastapi import HTTPException, Request
+from fastapi import Request
 from sqlalchemy.orm import Session
 
 from ee.onyx.server.gateway import api as gateway_api
 from onyx.db.models import User
+from onyx.db.usage import UsageLimitExceededError, UsageType
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
 from onyx.server import usage_limits
 from onyx.server.gateway.models import (
     AnthropicMessagesRequest,
@@ -72,7 +75,16 @@ def _provider(api_key: str) -> tuple[LLMProviderView, ModelConfigurationView]:
 
 
 def _over_cap(**_kwargs: object) -> None:
-    raise HTTPException(status_code=429, detail="LLM usage limit exceeded")
+    raise UsageLimitExceededError(usage_type=UsageType.LLM_COST, current=150, limit=100)
+
+
+@pytest.fixture(autouse=True)
+def _paid_tenant_with_cap() -> Generator[None, None, None]:
+    with (
+        patch.object(usage_limits, "is_tenant_on_trial_fn", return_value=False),
+        patch.object(usage_limits, "get_limit_for_usage_type", return_value=100),
+    ):
+        yield
 
 
 @pytest.mark.parametrize(
@@ -99,10 +111,10 @@ def test_managed_key_over_cap_never_reaches_the_provider(
         patch.object(usage_limits, "USAGE_LIMITS_ENABLED", True),
         patch.object(usage_limits, "_ONYX_MANAGED_API_KEYS", {_MANAGED_KEY}),
         patch.object(
-            usage_limits, "check_usage_and_raise", side_effect=_over_cap
+            usage_limits, "check_usage_limit", side_effect=_over_cap
         ) as usage_check,
         patch.multiple(gateway_api, **handlers),
-        pytest.raises(HTTPException) as exc_info,
+        pytest.raises(OnyxError) as exc_info,
     ):
         route(
             request=make_request(),
@@ -111,6 +123,7 @@ def test_managed_key_over_cap_never_reaches_the_provider(
             db_session=cast(Session, MagicMock(spec=Session)),
         )
 
+    assert exc_info.value.error_code is OnyxErrorCode.RATE_LIMITED
     assert exc_info.value.status_code == 429
     usage_check.assert_called_once()
     for handler in handlers.values():
@@ -133,7 +146,7 @@ def test_own_provider_key_is_not_subject_to_the_cap() -> None:
         patch.object(usage_limits, "USAGE_LIMITS_ENABLED", True),
         patch.object(usage_limits, "_ONYX_MANAGED_API_KEYS", {_MANAGED_KEY}),
         patch.object(
-            usage_limits, "check_usage_and_raise", side_effect=_over_cap
+            usage_limits, "check_usage_limit", side_effect=_over_cap
         ) as usage_check,
         patch.object(gateway_api, "handle_chat_completion") as handle,
     ):