mirror of
https://github.com/hanzoai/insights.git
synced 2026-11-02 14:32:45 +00:00
zeekay
1b367e2b62
The tree was a debranded snapshot of upstream taken around 2026-02-20, and upstream has moved about five months since. This replaces it with upstream master, mechanically re-derived: bin/debrand does the rename, and the pipeline and squashed migrations are restored from the shipped tree. Not a merge. The previous import replaced the tree wholesale in a single commit, so there is no shared content history to merge against. Reproducing the rename on the new tree is what actually converges. What survives: .hanzo/workflows, the Dockerfiles and bin/docker* entrypoints, compose files, LLM.md, and insights/migrations. The squashed initial migration is load-bearing — the live database records it by name, so upstream's ~918 individual migrations cannot come back. Nine products are dropped as unreferenced. Twenty-one more that we do not run are kept anyway because the core app still imports them; unpicking those is a refactor and does not belong in a rebase. This base does not build yet, and the reasons are not in the rename: upstream now wants Python 3.13.13 against our 3.12.12 and Django 5.2 against our 4.2, and the dependency manifests still name upstream packages that we republish under our own names. Those are decisions, not renames.
1528 lines
62 KiB
Python
1528 lines
62 KiB
Python
import pytest
|
|
|
|
from llm_gateway.auth.models import AuthenticatedUser
|
|
from llm_gateway.config import get_settings
|
|
from llm_gateway.rate_limiting.cost_throttles import CostThrottle
|
|
from llm_gateway.rate_limiting.throttles import ThrottleContext
|
|
|
|
|
|
def make_user(
|
|
user_id: int = 1, team_id: int = 1, auth_method: str = "oauth_access_token", is_staff: bool = False
|
|
) -> AuthenticatedUser:
|
|
return AuthenticatedUser(
|
|
user_id=user_id,
|
|
team_id=team_id,
|
|
auth_method=auth_method,
|
|
distinct_id=f"test-distinct-id-{user_id}",
|
|
scopes=["llm_gateway:read"],
|
|
is_staff=is_staff,
|
|
)
|
|
|
|
|
|
def make_context(
|
|
user: AuthenticatedUser | None = None,
|
|
product: str = "insights_code",
|
|
end_user_id: str | None = None,
|
|
plan_key: str | None = "insights-code-200-20260301",
|
|
seat_created_at: str | None = None,
|
|
seat_missing: bool = False,
|
|
code_usage_billed: bool = False,
|
|
billing_period_start: str | None = None,
|
|
) -> ThrottleContext:
|
|
user = user or make_user()
|
|
if end_user_id is None and user.auth_method == "oauth_access_token":
|
|
end_user_id = str(user.user_id)
|
|
return ThrottleContext(
|
|
user=user,
|
|
product=product,
|
|
end_user_id=end_user_id,
|
|
plan_key=plan_key,
|
|
seat_created_at=seat_created_at,
|
|
seat_missing=seat_missing,
|
|
code_usage_billed=code_usage_billed,
|
|
billing_period_start=billing_period_start,
|
|
)
|
|
|
|
|
|
async def recorded_cost(throttle: CostThrottle, context: ThrottleContext) -> float:
|
|
"""Cost the underlying limiter has actually accumulated for this context, read
|
|
straight from the limiter rather than through get_status (which reports staff as
|
|
unlimited). Lets a test prove spend was recorded even when enforcement is bypassed."""
|
|
return await throttle._get_limiter(context).get_current(throttle._get_cache_key(context))
|
|
|
|
|
|
class TestProductCostLimitConfig:
|
|
def test_default_product_cost_limits(self) -> None:
|
|
get_settings.cache_clear()
|
|
settings = get_settings()
|
|
assert "llm_gateway" in settings.product_cost_limits
|
|
assert settings.product_cost_limits["llm_gateway"].limit_usd == 1000.0
|
|
assert settings.product_cost_limits["llm_gateway"].window_seconds == 86400
|
|
assert settings.product_cost_limits["insights_ai"].limit_usd == 5000.0
|
|
assert settings.product_cost_limits["insights_ai"].window_seconds == 86400
|
|
|
|
def test_parses_json_string(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv(
|
|
"LLM_GATEWAY_PRODUCT_COST_LIMITS",
|
|
'{"wizard": {"limit_usd": 100, "window_seconds": 86400}, "insights_code": {"limit_usd": 50, "window_seconds": 14400}}',
|
|
)
|
|
get_settings.cache_clear()
|
|
settings = get_settings()
|
|
assert settings.product_cost_limits["wizard"].limit_usd == 100.0
|
|
assert settings.product_cost_limits["wizard"].window_seconds == 86400
|
|
assert settings.product_cost_limits["insights_code"].limit_usd == 50.0
|
|
assert settings.product_cost_limits["insights_code"].window_seconds == 14400
|
|
get_settings.cache_clear()
|
|
|
|
def test_legacy_twig_key_normalizes_to_insights_code(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv(
|
|
"LLM_GATEWAY_PRODUCT_COST_LIMITS",
|
|
'{"twig": {"limit_usd": 42, "window_seconds": 7200}}',
|
|
)
|
|
get_settings.cache_clear()
|
|
settings = get_settings()
|
|
assert "twig" not in settings.product_cost_limits
|
|
assert settings.product_cost_limits["insights_code"].limit_usd == 42.0
|
|
assert settings.product_cost_limits["insights_code"].window_seconds == 7200
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestUserCostLimitConfig:
|
|
def test_default_user_cost_limits(self) -> None:
|
|
get_settings.cache_clear()
|
|
settings = get_settings()
|
|
assert "insights_code" in settings.user_cost_limits
|
|
insights_code = settings.user_cost_limits["insights_code"]
|
|
assert insights_code.burst_limit_usd == 500.0
|
|
assert insights_code.burst_window_seconds == 86400
|
|
assert insights_code.sustained_limit_usd == 3000.0
|
|
assert insights_code.sustained_window_seconds == 2592000
|
|
get_settings.cache_clear()
|
|
|
|
def test_parses_json_string(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv(
|
|
"LLM_GATEWAY_USER_COST_LIMITS",
|
|
'{"insights_code": {"burst_limit_usd": 50, "burst_window_seconds": 86400, "sustained_limit_usd": 500, "sustained_window_seconds": 2592000}}',
|
|
)
|
|
get_settings.cache_clear()
|
|
settings = get_settings()
|
|
insights_code = settings.user_cost_limits["insights_code"]
|
|
assert insights_code.burst_limit_usd == 50.0
|
|
assert insights_code.sustained_limit_usd == 500.0
|
|
get_settings.cache_clear()
|
|
|
|
def test_unset_env_returns_defaults(self) -> None:
|
|
get_settings.cache_clear()
|
|
settings = get_settings()
|
|
assert "insights_code" in settings.user_cost_limits
|
|
get_settings.cache_clear()
|
|
|
|
def test_legacy_twig_key_normalizes_to_insights_code(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv(
|
|
"LLM_GATEWAY_USER_COST_LIMITS",
|
|
'{"twig": {"burst_limit_usd": 77, "burst_window_seconds": 3600, "sustained_limit_usd": 777, "sustained_window_seconds": 86400}}',
|
|
)
|
|
get_settings.cache_clear()
|
|
settings = get_settings()
|
|
assert "twig" not in settings.user_cost_limits
|
|
assert settings.user_cost_limits["insights_code"].burst_limit_usd == 77.0
|
|
assert settings.user_cost_limits["insights_code"].sustained_limit_usd == 777.0
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestProductCostThrottle:
|
|
@pytest.mark.asyncio
|
|
async def test_allows_when_under_limit(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
context = make_context(product="llm_gateway")
|
|
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_denies_when_over_limit(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
context = make_context(product="llm_gateway")
|
|
|
|
await throttle.record_cost(context, 1000.0)
|
|
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
assert result.scope == "product_cost"
|
|
assert result.detail == "Product rate limit exceeded"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_products_have_separate_limits(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
|
|
ctx_wizard = make_context(product="wizard")
|
|
ctx_insights_code = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(ctx_wizard, 10000.0)
|
|
|
|
result_wizard = await throttle.allow_request(ctx_wizard)
|
|
result_insights_code = await throttle.allow_request(ctx_insights_code)
|
|
|
|
assert result_wizard.allowed is False
|
|
assert result_insights_code.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_format(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
context = make_context(product="wizard")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:product:wizard"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_status_for_product_reports_used_and_limit(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
context = make_context(product="llm_gateway")
|
|
await throttle.record_cost(context, 42.0)
|
|
|
|
status = await throttle.get_status_for_product("llm_gateway")
|
|
|
|
assert status is not None
|
|
assert status.used_usd == pytest.approx(42.0)
|
|
assert status.limit_usd == 1000.0
|
|
assert status.remaining_usd == pytest.approx(958.0)
|
|
assert status.exceeded is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_status_for_product_returns_none_for_unknown(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
|
|
assert await throttle.get_status_for_product("not_a_real_product") is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_status_for_product_ignores_staff_multiplier_suffix(
|
|
self, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""Gauge readings track only the shared (multiplier=1) pool — staff spend lands in a
|
|
suffixed Redis bucket and is intentionally invisible to the gauge, so alerts don't
|
|
double-count staff against the shared cap."""
|
|
monkeypatch.setenv("LLM_GATEWAY_STAFF_RATE_LIMIT_MULTIPLIER", "10")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
|
|
staff_context = make_context(user=make_user(user_id=1, is_staff=True), product="llm_gateway")
|
|
await throttle.record_cost(staff_context, 50.0)
|
|
|
|
status = await throttle.get_status_for_product("llm_gateway")
|
|
assert status is not None
|
|
assert status.used_usd == pytest.approx(0.0), "staff spend must not appear in the shared-pool gauge"
|
|
|
|
shared_context = make_context(user=make_user(user_id=2, is_staff=False), product="llm_gateway")
|
|
await throttle.record_cost(shared_context, 7.0)
|
|
|
|
status = await throttle.get_status_for_product("llm_gateway")
|
|
assert status is not None
|
|
assert status.used_usd == pytest.approx(7.0), "gauge should reflect shared-pool spend only"
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestUserCostBurstThrottle:
|
|
@pytest.mark.asyncio
|
|
async def test_allows_when_under_limit(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_denies_when_over_burst_limit(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 500.0)
|
|
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
assert result.scope == "user_cost_burst"
|
|
assert result.detail == "User burst rate limit exceeded"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uses_default_for_products_without_config(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="llm_gateway")
|
|
|
|
await throttle.record_cost(context, 99.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
await throttle.record_cost(context, 1.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skips_without_end_user_id(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(user_id=1, auth_method="personal_api_key")
|
|
context = make_context(user=user, product="insights_code", end_user_id=None)
|
|
|
|
await throttle.record_cost(context, 99999.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_includes_product_and_scope(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code", end_user_id="42")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_burst:insights_code:42"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_users_have_separate_limits(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
|
|
user1 = make_user(user_id=1)
|
|
user2 = make_user(user_id=2)
|
|
ctx1 = make_context(user=user1, product="insights_code")
|
|
ctx2 = make_context(user=user2, product="insights_code")
|
|
|
|
await throttle.record_cost(ctx1, 500.0)
|
|
|
|
assert (await throttle.allow_request(ctx1)).allowed is False
|
|
assert (await throttle.allow_request(ctx2)).allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestUserCostSustainedThrottle:
|
|
@pytest.mark.asyncio
|
|
async def test_allows_when_under_limit(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_denies_when_over_sustained_limit(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 3000.0)
|
|
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
assert result.scope == "user_cost_sustained"
|
|
assert result.detail == "User sustained rate limit exceeded"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uses_default_for_products_without_config(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="llm_gateway")
|
|
|
|
await throttle.record_cost(context, 999.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
await throttle.record_cost(context, 1.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_includes_product_scope_and_period(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code", end_user_id="42")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_sustained:insights_code:42:period:0"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_includes_period_for_free_plan(self) -> None:
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
created = (datetime.now(tz=UTC) - timedelta(days=5)).isoformat()
|
|
context = make_context(
|
|
product="insights_code",
|
|
end_user_id="42",
|
|
plan_key="insights-code-free-20260301",
|
|
seat_created_at=created,
|
|
)
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_sustained:insights_code:42:period:0"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_period_increments_after_period_days(self) -> None:
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
created = (datetime.now(tz=UTC) - timedelta(days=35)).isoformat()
|
|
context = make_context(
|
|
product="insights_code",
|
|
end_user_id="42",
|
|
plan_key="insights-code-free-20260301",
|
|
seat_created_at=created,
|
|
)
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_sustained:insights_code:42:period:1"
|
|
|
|
def test_cache_key_uses_billing_period_start_over_seat_created_at(self) -> None:
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
old_seat = (datetime.now(tz=UTC) - timedelta(days=35)).isoformat()
|
|
recent_period = (datetime.now(tz=UTC) - timedelta(days=5)).isoformat()
|
|
context = make_context(
|
|
product="insights_code",
|
|
end_user_id="42",
|
|
plan_key="insights-code-pro-200-20260301",
|
|
seat_created_at=old_seat,
|
|
billing_period_start=recent_period,
|
|
)
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_sustained:insights_code:42:period:0"
|
|
|
|
def test_cache_key_falls_back_to_seat_created_at_without_billing_period(self) -> None:
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
old_seat = (datetime.now(tz=UTC) - timedelta(days=35)).isoformat()
|
|
context = make_context(
|
|
product="insights_code",
|
|
end_user_id="42",
|
|
plan_key="insights-code-pro-200-20260301",
|
|
seat_created_at=old_seat,
|
|
billing_period_start=None,
|
|
)
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_sustained:insights_code:42:period:1"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_includes_period_for_pro_plan(self) -> None:
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
created = (datetime.now(tz=UTC) - timedelta(days=5)).isoformat()
|
|
context = make_context(
|
|
product="insights_code",
|
|
end_user_id="42",
|
|
plan_key="insights-code-200-20260301",
|
|
seat_created_at=created,
|
|
)
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_sustained:insights_code:42:period:0"
|
|
|
|
|
|
class TestBurstSustainedInteraction:
|
|
@pytest.mark.asyncio
|
|
async def test_burst_denies_before_sustained(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle, UserCostSustainedThrottle
|
|
|
|
burst = UserCostBurstThrottle(redis=None)
|
|
sustained = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await burst.record_cost(context, 500.0)
|
|
await sustained.record_cost(context, 100.0)
|
|
|
|
assert (await burst.allow_request(context)).allowed is False
|
|
assert (await sustained.allow_request(context)).allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sustained_denies_even_if_burst_allows(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle, UserCostSustainedThrottle
|
|
|
|
burst = UserCostBurstThrottle(redis=None)
|
|
sustained = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await burst.record_cost(context, 50.0)
|
|
await sustained.record_cost(context, 3000.0)
|
|
|
|
assert (await burst.allow_request(context)).allowed is True
|
|
assert (await sustained.allow_request(context)).allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_limits_via_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv(
|
|
"LLM_GATEWAY_USER_COST_LIMITS",
|
|
'{"insights_code": {"burst_limit_usd": 200, "burst_window_seconds": 86400, "sustained_limit_usd": 2000, "sustained_window_seconds": 2592000}}',
|
|
)
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 150.0)
|
|
assert (await throttle.allow_request(context)).allowed is True
|
|
|
|
await throttle.record_cost(context, 60.0)
|
|
assert (await throttle.allow_request(context)).allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestUserCostDisabledFlag:
|
|
@pytest.mark.asyncio
|
|
async def test_allows_when_limits_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_USER_COST_LIMITS_DISABLED", "true")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 1000.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tracks_but_does_not_enforce_when_disabled(
|
|
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
|
) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_USER_COST_LIMITS_DISABLED", "true")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 1000.0)
|
|
result = await throttle.allow_request(context)
|
|
|
|
captured = capsys.readouterr()
|
|
assert result.allowed is True
|
|
assert "cost_throttle_exceeded" in captured.out
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enforces_when_limits_enabled(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 1000.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
assert result.scope == "user_cost_burst"
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestStaffUnlimitedUsage:
|
|
@pytest.mark.asyncio
|
|
async def test_staff_allowed_over_burst_limit(self) -> None:
|
|
# staff_unlimited_usage defaults to True.
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(user=make_user(is_staff=True), product="insights_code")
|
|
|
|
await throttle.record_cost(context, 100_000.0)
|
|
# Observability guarantee: staff spend is still recorded even though the
|
|
# request is allowed. Fails if record_cost starts skipping staff.
|
|
assert await recorded_cost(throttle, context) == 100_000.0
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_staff_allowed_over_sustained_limit(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(user=make_user(is_staff=True), product="insights_code")
|
|
|
|
await throttle.record_cost(context, 100_000.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_staff_status_reports_unlimited(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(user=make_user(is_staff=True), product="insights_code")
|
|
|
|
await throttle.record_cost(context, 100_000.0)
|
|
# Spend is recorded on the limiter even though the reported status hides it
|
|
# behind an unlimited budget — the observability guarantee this PR promises.
|
|
assert await recorded_cost(throttle, context) == 100_000.0
|
|
status = await throttle.get_status(context)
|
|
assert status.used_usd == 0.0
|
|
assert status.exceeded is False
|
|
assert status.limit_usd == float("inf")
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_free_plan_staff_still_unlimited(self) -> None:
|
|
# The case that bit staff before: a staff member on the free plan was
|
|
# pinned to the (multiplied) free-plan cap rather than treated as unlimited.
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(
|
|
user=make_user(is_staff=True),
|
|
product="insights_code",
|
|
plan_key="insights-code-free-20260301",
|
|
seat_created_at=(datetime.now(tz=UTC) - timedelta(days=5)).isoformat(),
|
|
)
|
|
|
|
await throttle.record_cost(context, 100_000.0)
|
|
assert (await throttle.allow_request(context)).allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_staff_still_enforced(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(user=make_user(is_staff=False), product="insights_code")
|
|
|
|
await throttle.record_cost(context, 500.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_staff_enforced_when_setting_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_STAFF_UNLIMITED_USAGE", "false")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
# With the bypass off, staff fall back to the multiplied finite cap
|
|
# ($500 burst * 10 staff multiplier = $5000).
|
|
context = make_context(user=make_user(is_staff=True), product="insights_code")
|
|
|
|
await throttle.record_cost(context, 5_000.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestRetryAfterHeader:
|
|
@pytest.mark.asyncio
|
|
async def test_retry_after_returns_full_window_without_redis(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 500.0)
|
|
result = await throttle.allow_request(context)
|
|
|
|
assert result.allowed is False
|
|
assert result.retry_after == 86400
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry_after_returns_ttl_from_redis(self) -> None:
|
|
get_settings.cache_clear()
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
mock_redis = MagicMock()
|
|
mock_redis.get = AsyncMock(return_value=b"1000.0")
|
|
mock_redis.ttl = AsyncMock(return_value=600)
|
|
|
|
throttle = UserCostBurstThrottle(redis=mock_redis)
|
|
context = make_context(product="insights_code")
|
|
|
|
result = await throttle.allow_request(context)
|
|
|
|
assert result.allowed is False
|
|
assert result.retry_after == 600
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestCostAccumulation:
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_small_costs_accumulate_to_limit(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
for _ in range(49):
|
|
await throttle.record_cost(context, 10.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
await throttle.record_cost(context, 20.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zero_cost_not_recorded(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 0.0)
|
|
await throttle.record_cost(context, -1.0)
|
|
|
|
limiter = throttle._get_limiter(context)
|
|
key = throttle._get_cache_key(context)
|
|
current = await limiter.get_current(key)
|
|
assert current == 0.0
|
|
|
|
|
|
class TestCostRateLimiterRedisIntegration:
|
|
@pytest.mark.asyncio
|
|
async def test_redis_incr_called_with_correct_args(self) -> None:
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
mock_redis = MagicMock()
|
|
mock_redis.eval = AsyncMock(return_value=0.5)
|
|
mock_redis.get = AsyncMock(return_value=b"0.0")
|
|
|
|
throttle = UserCostBurstThrottle(redis=mock_redis)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 0.5)
|
|
|
|
mock_redis.eval.assert_called_once()
|
|
call_args = mock_redis.eval.call_args
|
|
assert "ratelimit:cost:user:user_cost_burst:insights_code:1" in call_args[0]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_get_current_returns_accumulated_cost(self) -> None:
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
mock_redis = MagicMock()
|
|
mock_redis.get = AsyncMock(return_value=b"1.5")
|
|
|
|
throttle = UserCostBurstThrottle(redis=mock_redis)
|
|
context = make_context(product="insights_code")
|
|
|
|
limiter = throttle._get_limiter(context)
|
|
current = await limiter.get_current(throttle._get_cache_key(context))
|
|
|
|
assert current == 1.5
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_ttl_returns_remaining_time(self) -> None:
|
|
get_settings.cache_clear()
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
mock_redis = MagicMock()
|
|
mock_redis.get = AsyncMock(return_value=b"1000.0")
|
|
mock_redis.ttl = AsyncMock(return_value=1800)
|
|
|
|
throttle = UserCostBurstThrottle(redis=mock_redis)
|
|
context = make_context(product="insights_code")
|
|
|
|
result = await throttle.allow_request(context)
|
|
|
|
assert result.allowed is False
|
|
assert result.retry_after == 1800
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_falls_back_to_local_on_redis_error(self) -> None:
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
mock_redis = MagicMock()
|
|
mock_redis.eval = AsyncMock(side_effect=Exception("Redis error"))
|
|
mock_redis.get = AsyncMock(side_effect=Exception("Redis error"))
|
|
|
|
throttle = UserCostBurstThrottle(redis=mock_redis)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 0.1)
|
|
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
|
|
class TestRateLimitMultiplier:
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_has_no_suffix_for_plain_user(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_TEAM_RATE_LIMIT_MULTIPLIERS", "{}")
|
|
monkeypatch.setenv("LLM_GATEWAY_STAFF_RATE_LIMIT_MULTIPLIER", "10")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(user_id=1, team_id=99, is_staff=False)
|
|
context = make_context(user=user, product="insights_code")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert ":m" not in key
|
|
assert key == "cost:user:user_cost_burst:insights_code:1"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_suffix_from_team_multiplier(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_TEAM_RATE_LIMIT_MULTIPLIERS", '{"2": 10}')
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(user_id=1, team_id=2, is_staff=False)
|
|
context = make_context(user=user, product="insights_code")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_burst:insights_code:1:m10"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_suffix_from_staff_multiplier(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_STAFF_RATE_LIMIT_MULTIPLIER", "10")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
# Staff on an unconfigured team still gets the suffixed bucket.
|
|
user = make_user(user_id=1, team_id=99, is_staff=True)
|
|
context = make_context(user=user, product="insights_code")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_burst:insights_code:1:m10"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_gets_higher_limit(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_TEAM_RATE_LIMIT_MULTIPLIERS", '{"2": 10}')
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(user_id=1, team_id=2, is_staff=False)
|
|
context = make_context(user=user, product="insights_code")
|
|
|
|
await throttle.record_cost(context, 100.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True, "Should allow - team has 10x multiplier ($2000 limit vs $100 used)"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_staff_gets_higher_limit_on_any_team(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_STAFF_RATE_LIMIT_MULTIPLIER", "10")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
# Staff on an arbitrary team — the impersonation case — still gets the elevated cap.
|
|
user = make_user(user_id=1, team_id=99, is_staff=True)
|
|
context = make_context(user=user, product="insights_code")
|
|
|
|
await throttle.record_cost(context, 100.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True, "Should allow - staff has 10x multiplier ($2000 limit vs $100 used)"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_effective_multiplier_is_max_of_team_and_staff(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_TEAM_RATE_LIMIT_MULTIPLIERS", '{"2": 10}')
|
|
monkeypatch.setenv("LLM_GATEWAY_STAFF_RATE_LIMIT_MULTIPLIER", "3")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
# Staff on a configured team gets the larger of the two multipliers (10, not 3).
|
|
user = make_user(user_id=1, team_id=2, is_staff=True)
|
|
context = make_context(user=user, product="insights_code")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:user:user_cost_burst:insights_code:1:m10"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_product_cache_key_includes_multiplier_suffix(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_STAFF_RATE_LIMIT_MULTIPLIER", "10")
|
|
get_settings.cache_clear()
|
|
|
|
from llm_gateway.rate_limiting.cost_throttles import ProductCostThrottle
|
|
|
|
throttle = ProductCostThrottle(redis=None)
|
|
user = make_user(user_id=1, is_staff=True)
|
|
context = make_context(user=user, product="wizard")
|
|
|
|
key = throttle._get_cache_key(context)
|
|
assert key == "cost:product:wizard:m10"
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestUnconfiguredProductsUseDefaults:
|
|
"""Products without user_cost_limits config use default limits ($100/24h burst, $1000/30d sustained)."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unconfigured_product_uses_burst_default(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="wizard")
|
|
|
|
await throttle.record_cost(context, 99.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
await throttle.record_cost(context, 1.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unconfigured_product_uses_sustained_default(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="unconfigured_product")
|
|
|
|
await throttle.record_cost(context, 999.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
await throttle.record_cost(context, 1.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_configured_and_unconfigured_products_both_limited(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle, UserCostSustainedThrottle
|
|
|
|
burst = UserCostBurstThrottle(redis=None)
|
|
sustained = UserCostSustainedThrottle(redis=None)
|
|
|
|
ctx_insights_code = make_context(product="insights_code")
|
|
ctx_wizard = make_context(product="wizard")
|
|
|
|
await burst.record_cost(ctx_insights_code, 500.0)
|
|
await burst.record_cost(ctx_wizard, 100.0)
|
|
await sustained.record_cost(ctx_insights_code, 3000.0)
|
|
await sustained.record_cost(ctx_wizard, 1000.0)
|
|
|
|
assert (await burst.allow_request(ctx_insights_code)).allowed is False
|
|
assert (await burst.allow_request(ctx_wizard)).allowed is False
|
|
assert (await sustained.allow_request(ctx_insights_code)).allowed is False
|
|
assert (await sustained.allow_request(ctx_wizard)).allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_logs_info_for_unconfigured_product_with_end_user(self, capsys: pytest.CaptureFixture[str]) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle, _UserCostThrottleBase
|
|
|
|
_UserCostThrottleBase._warned_products = set()
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="unconfigured_product")
|
|
|
|
await throttle.allow_request(context)
|
|
captured = capsys.readouterr()
|
|
assert "user_cost_limits_using_default" in captured.out
|
|
assert "unconfigured_product" in captured.out
|
|
|
|
await throttle.allow_request(context)
|
|
captured2 = capsys.readouterr()
|
|
assert "user_cost_limits_using_default" not in captured2.out
|
|
|
|
_UserCostThrottleBase._warned_products = set()
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_warning_for_unconfigured_product_without_end_user(
|
|
self, capsys: pytest.CaptureFixture[str]
|
|
) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle, _UserCostThrottleBase
|
|
|
|
_UserCostThrottleBase._warned_products = set()
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(auth_method="personal_api_key")
|
|
context = make_context(user=user, product="wizard", end_user_id=None)
|
|
|
|
await throttle.allow_request(context)
|
|
captured = capsys.readouterr()
|
|
assert "user_cost_limits_using_default" not in captured.out
|
|
|
|
_UserCostThrottleBase._warned_products = set()
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dynamically_adding_product_config_overrides_defaults(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="wizard")
|
|
|
|
await throttle.record_cost(context, 99.0)
|
|
assert (await throttle.allow_request(context)).allowed is True
|
|
|
|
monkeypatch.setenv(
|
|
"LLM_GATEWAY_USER_COST_LIMITS",
|
|
'{"insights_code": {"burst_limit_usd": 100, "burst_window_seconds": 86400, "sustained_limit_usd": 1000, "sustained_window_seconds": 2592000}, '
|
|
'"wizard": {"burst_limit_usd": 50, "burst_window_seconds": 3600, "sustained_limit_usd": 200, "sustained_window_seconds": 86400}}',
|
|
)
|
|
get_settings.cache_clear()
|
|
|
|
await throttle.record_cost(context, 50.0)
|
|
assert (await throttle.allow_request(context)).allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestUserCostEdgeCases:
|
|
@pytest.mark.asyncio
|
|
async def test_sustained_skips_without_end_user_id(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
user = make_user(user_id=1, auth_method="personal_api_key")
|
|
context = make_context(user=user, product="insights_code", end_user_id=None)
|
|
|
|
await throttle.record_cost(context, 99999.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_string_end_user_id_treated_as_no_user(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(user_id=1, auth_method="personal_api_key")
|
|
context = make_context(user=user, product="insights_code", end_user_id="")
|
|
|
|
await throttle.record_cost(context, 99999.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_personal_api_key_with_end_user_id_enforces_limits(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(user_id=1, auth_method="personal_api_key")
|
|
context = make_context(user=user, product="insights_code", end_user_id="ext-user-42")
|
|
|
|
await throttle.record_cost(context, 500.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
assert result.scope == "user_cost_burst"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_config_returns_default_for_unconfigured_product(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="unconfigured_product")
|
|
|
|
config = throttle._get_config(context)
|
|
assert config.burst_limit_usd == 100.0
|
|
assert config.burst_window_seconds == 86400
|
|
assert config.sustained_limit_usd == 1000.0
|
|
assert config.sustained_window_seconds == 2592000
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_key_empty_without_end_user_id(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(auth_method="personal_api_key")
|
|
context = make_context(user=user, product="insights_code", end_user_id=None)
|
|
|
|
assert throttle._get_cache_key(context) == ""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_products_same_user_isolated(self) -> None:
|
|
get_settings.cache_clear()
|
|
monkeypatch = pytest.MonkeyPatch()
|
|
monkeypatch.setenv(
|
|
"LLM_GATEWAY_USER_COST_LIMITS",
|
|
'{"insights_code": {"burst_limit_usd": 100, "burst_window_seconds": 86400, "sustained_limit_usd": 1000, "sustained_window_seconds": 2592000}, '
|
|
'"wizard": {"burst_limit_usd": 50, "burst_window_seconds": 3600, "sustained_limit_usd": 200, "sustained_window_seconds": 86400}}',
|
|
)
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
|
|
ctx_insights_code = make_context(product="insights_code", end_user_id="42")
|
|
ctx_wizard = make_context(product="wizard", end_user_id="42")
|
|
|
|
await throttle.record_cost(ctx_wizard, 50.0)
|
|
|
|
assert (await throttle.allow_request(ctx_wizard)).allowed is False
|
|
assert (await throttle.allow_request(ctx_insights_code)).allowed is True
|
|
monkeypatch.undo()
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_just_below_limit_still_allowed(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 199.99)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_exactly_at_limit_denied(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 500.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestUserCostDisabledSustained:
|
|
@pytest.mark.asyncio
|
|
async def test_sustained_allows_when_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_USER_COST_LIMITS_DISABLED", "true")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await throttle.record_cost(context, 9999.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_both_burst_and_sustained_allow_when_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LLM_GATEWAY_USER_COST_LIMITS_DISABLED", "true")
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle, UserCostSustainedThrottle
|
|
|
|
burst = UserCostBurstThrottle(redis=None)
|
|
sustained = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code")
|
|
|
|
await burst.record_cost(context, 9999.0)
|
|
await sustained.record_cost(context, 9999.0)
|
|
|
|
assert (await burst.allow_request(context)).allowed is True
|
|
assert (await sustained.allow_request(context)).allowed is True
|
|
get_settings.cache_clear()
|
|
|
|
|
|
class TestRateLimitPoisoningPrevention:
|
|
"""Verify that cost is always recorded against the authenticated user's own bucket.
|
|
|
|
An attacker must not be able to:
|
|
1. Poison a victim's rate limit bucket by injecting their user ID
|
|
2. Bypass rate limiting entirely by omitting user identification
|
|
"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_recorded_against_own_bucket_not_victim(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
|
|
attacker = make_user(user_id=999, auth_method="personal_api_key")
|
|
victim = make_user(user_id=42, auth_method="oauth_access_token")
|
|
|
|
attacker_ctx = make_context(user=attacker, product="insights_code", end_user_id="999")
|
|
victim_ctx = make_context(user=victim, product="insights_code")
|
|
|
|
await throttle.record_cost(attacker_ctx, 500.0)
|
|
|
|
assert (await throttle.allow_request(attacker_ctx)).allowed is False, (
|
|
"Attacker's own bucket should be exhausted"
|
|
)
|
|
assert (await throttle.allow_request(victim_ctx)).allowed is True, (
|
|
"Victim's bucket must not be affected by attacker's usage"
|
|
)
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_key_user_with_end_user_id_is_rate_limited(self) -> None:
|
|
get_settings.cache_clear()
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
user = make_user(user_id=999, auth_method="personal_api_key")
|
|
context = make_context(user=user, product="insights_code", end_user_id="999")
|
|
|
|
await throttle.record_cost(context, 500.0)
|
|
result = await throttle.allow_request(context)
|
|
|
|
assert result.allowed is False, "Personal API key users must be rate limited when end_user_id is set"
|
|
get_settings.cache_clear()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_injected_victim_id_does_not_affect_victim_cache_key(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
|
|
attacker = make_user(user_id=999, auth_method="personal_api_key")
|
|
attacker_ctx = make_context(user=attacker, product="insights_code", end_user_id="999")
|
|
attacker_key = throttle._get_cache_key(attacker_ctx)
|
|
|
|
victim = make_user(user_id=42, auth_method="oauth_access_token")
|
|
victim_ctx = make_context(user=victim, product="insights_code")
|
|
victim_key = throttle._get_cache_key(victim_ctx)
|
|
|
|
assert attacker_key != victim_key
|
|
assert ":999" in attacker_key
|
|
assert ":42" in victim_key
|
|
|
|
|
|
class TestPlanAwareThrottling:
|
|
@pytest.mark.asyncio
|
|
async def test_free_user_gets_limit_of_20(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key=None, seat_created_at="2026-01-01T00:00:00+00:00")
|
|
|
|
await throttle.record_cost(context, 20.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_free_user_gets_sustained_limit_of_20(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key=None, seat_created_at="2026-01-01T00:00:00+00:00")
|
|
|
|
await throttle.record_cost(context, 20.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_old_free_user_still_gets_burst_limit(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key=None, seat_created_at="2025-01-01T00:00:00+00:00")
|
|
|
|
await throttle.record_cost(context, 4.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_old_free_user_still_gets_sustained_limit(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostSustainedThrottle
|
|
|
|
throttle = UserCostSustainedThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key=None, seat_created_at="2025-01-01T00:00:00+00:00")
|
|
|
|
await throttle.record_cost(context, 19.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_seat_gets_pro_limits(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key=None, seat_created_at=None)
|
|
|
|
await throttle.record_cost(context, 50.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pro_plan_allows_higher_usage(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key="insights-code-200-20260301")
|
|
|
|
await throttle.record_cost(context, 50.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_renamed_pro_plan_allows_higher_usage(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key="insights-code-pro-200-20260301")
|
|
|
|
await throttle.record_cost(context, 50.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zero_dollar_pro_plan_allows_higher_usage(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="insights_code", plan_key="insights-code-pro-0-20260422")
|
|
|
|
await throttle.record_cost(context, 50.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_free_plan_key_gets_free_limits(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(
|
|
product="insights_code",
|
|
plan_key="insights-code-free-20260301",
|
|
seat_created_at="2026-01-01T00:00:00+00:00",
|
|
)
|
|
|
|
await throttle.record_cost(context, 4.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("seat_missing", "expected_allowed"),
|
|
[
|
|
# Confirmed no seat (404) → free cap: $80 spent exceeds the $20 limit.
|
|
(True, False),
|
|
# Plan resolution failed → default limits ($500/day): fail loose so an
|
|
# outage never clamps paying users whose plan couldn't be resolved.
|
|
(False, True),
|
|
],
|
|
)
|
|
async def test_seatless_user_limit_depends_on_seat_missing(
|
|
self, seat_missing: bool, expected_allowed: bool
|
|
) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(
|
|
product="insights_code",
|
|
plan_key=None,
|
|
seat_created_at=None,
|
|
seat_missing=seat_missing,
|
|
)
|
|
|
|
await throttle.record_cost(context, 80.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is expected_allowed
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("seat_missing", "code_usage_billed", "expected_allowed"),
|
|
[
|
|
# Seatless + org-billed Desktop usage: bills to the org - no per-user cap.
|
|
# $600 spent clears both the $20 free cap and the $500/day default.
|
|
(True, True, True),
|
|
# Seatless + not billed: the free cap holds (this pair pins that the
|
|
# bypass never leaks to orgs that would get the usage for free).
|
|
(True, False, False),
|
|
# Seat still exists (pre-retirement): org billing doesn't lift the cap -
|
|
# seat-covered usage stays bounded until the seat is deleted.
|
|
(False, True, False),
|
|
],
|
|
)
|
|
async def test_code_usage_billing_gates_seatless_cap_bypass(
|
|
self, seat_missing: bool, code_usage_billed: bool, expected_allowed: bool
|
|
) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(
|
|
product="insights_code",
|
|
plan_key=None if seat_missing else "insights-code-free-20260301",
|
|
seat_created_at=None if seat_missing else "2026-01-01T00:00:00+00:00",
|
|
seat_missing=seat_missing,
|
|
code_usage_billed=code_usage_billed,
|
|
)
|
|
|
|
await throttle.record_cost(context, 600.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is expected_allowed
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_org_billed_seatless_status_reports_unlimited(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(
|
|
product="insights_code", plan_key=None, seat_created_at=None, seat_missing=True, code_usage_billed=True
|
|
)
|
|
|
|
await throttle.record_cost(context, 600.0)
|
|
status = await throttle.get_status(context)
|
|
# The usage endpoint renders this - it must agree with enforcement (never
|
|
# show a paying user as rate limited when nothing would block them).
|
|
assert status.exceeded is False
|
|
assert status.limit_usd == float("inf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_code_product_ignores_plan(self) -> None:
|
|
from llm_gateway.rate_limiting.cost_throttles import UserCostBurstThrottle
|
|
|
|
throttle = UserCostBurstThrottle(redis=None)
|
|
context = make_context(product="wizard", plan_key=None)
|
|
|
|
await throttle.record_cost(context, 50.0)
|
|
result = await throttle.allow_request(context)
|
|
assert result.allowed is True
|
|
|
|
|
|
class TestCostAccumulatorTTL:
|
|
def test_cost_expires_after_window(self) -> None:
|
|
from unittest.mock import patch
|
|
|
|
from llm_gateway.rate_limiting.redis_limiter import CostAccumulator
|
|
|
|
accumulator = CostAccumulator(limit=10.0, window_seconds=60)
|
|
|
|
with patch("llm_gateway.rate_limiting.redis_limiter.time.monotonic", return_value=0):
|
|
accumulator.incr("user1", 5.0)
|
|
assert accumulator.get_current("user1") == 5.0
|
|
|
|
with patch("llm_gateway.rate_limiting.redis_limiter.time.monotonic", return_value=61):
|
|
assert accumulator.get_current("user1") == 0.0
|
|
|
|
def test_cost_persists_within_window(self) -> None:
|
|
from unittest.mock import patch
|
|
|
|
from llm_gateway.rate_limiting.redis_limiter import CostAccumulator
|
|
|
|
accumulator = CostAccumulator(limit=10.0, window_seconds=60)
|
|
|
|
with patch("llm_gateway.rate_limiting.redis_limiter.time.monotonic", return_value=0):
|
|
accumulator.incr("user1", 5.0)
|
|
|
|
with patch("llm_gateway.rate_limiting.redis_limiter.time.monotonic", return_value=30):
|
|
assert accumulator.get_current("user1") == 5.0
|
|
accumulator.incr("user1", 3.0)
|
|
assert accumulator.get_current("user1") == 8.0
|
|
|
|
def test_new_window_starts_fresh(self) -> None:
|
|
from unittest.mock import patch
|
|
|
|
from llm_gateway.rate_limiting.redis_limiter import CostAccumulator
|
|
|
|
accumulator = CostAccumulator(limit=10.0, window_seconds=60)
|
|
|
|
with patch("llm_gateway.rate_limiting.redis_limiter.time.monotonic", return_value=0):
|
|
accumulator.incr("user1", 10.0)
|
|
assert accumulator.incr("user1", 1.0) is False
|
|
|
|
with patch("llm_gateway.rate_limiting.redis_limiter.time.monotonic", return_value=61):
|
|
assert accumulator.incr("user1", 5.0) is True
|
|
assert accumulator.get_current("user1") == 5.0
|
|
|
|
def test_different_keys_independent(self) -> None:
|
|
from llm_gateway.rate_limiting.redis_limiter import CostAccumulator
|
|
|
|
accumulator = CostAccumulator(limit=10.0, window_seconds=60)
|
|
|
|
accumulator.incr("user1", 5.0)
|
|
accumulator.incr("user2", 3.0)
|
|
|
|
assert accumulator.get_current("user1") == 5.0
|
|
assert accumulator.get_current("user2") == 3.0
|