mirror of
https://github.com/hanzoai/insights.git
synced 2026-10-30 16:15:31 +00:00
zeekay
19fdf55ceb
A key now opens with what it is allowed to touch. `pk-` is publishable and safe in a browser; `sk-` is secret and never leaves a server. The five marks it replaces all began `ph` -- the upstream's initials -- and they were the last of its brand a customer read back to us, in every SDK snippet. No word rule could reach them, because `ph` is not a word. The marks were also not telling the truth. The constant said `phc_` while every live project token said `hi_`, so the middleware that routes on the constant matched no real token at all. Two spellings existed and one was fiction. Minting now happens in exactly one function and reading in exactly one, so a second spelling has nowhere to live. `phx_` (personal) and `phs_` (project secret) both become `sk-`. The mark now says only that a value is secret; which store holds it is settled by asking. That was already how personal keys worked -- keys minted before the marks carry none, so the store always decided -- and both secret resolvers already returned None on a miss and fell through, so one shared extractor replaces the two that could drift apart. OAuth's access and refresh tokens stay distinguishable as `at-` and `rt-`, because the parser tells them apart before either is looked up, and the domain-verification token becomes `vt-` rather than keeping a sixth mint of its own. `bin/debrand --check` now refuses the marks as a shape, so a seventh cannot be coined. Two rules that were about to rename things that are not ours to spell are bounded: `Lemon` no longer eats `LemonSqueezy`, and the `hog_flow` rules no longer reach a Postgres table, a Datastore table or a route. Keys are stored hashed, so a mark cannot be rewritten in place: the three live project tokens have to be reminted.
309 lines
11 KiB
Python
309 lines
11 KiB
Python
import time
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fakeredis import aioredis as fakeredis
|
|
from fastapi.testclient import TestClient
|
|
|
|
from llm_gateway.rate_limiting.redis_limiter import RateLimiter
|
|
from llm_gateway.rate_limiting.throttles import Throttle, ThrottleContext, ThrottleResult
|
|
from llm_gateway.rate_limiting.token_bucket import TokenBucketLimiter
|
|
from tests.conftest import create_test_app
|
|
|
|
|
|
class TestTokenBucketLimiter:
|
|
@pytest.mark.parametrize(
|
|
"rate,capacity,consume_count,expected_results",
|
|
[
|
|
pytest.param(0.0, 3.0, 5, [True, True, True, False, False], id="exhaust_bucket_capacity_3"),
|
|
pytest.param(0.0, 1.0, 3, [True, False, False], id="exhaust_bucket_capacity_1"),
|
|
pytest.param(1000.0, 100.0, 5, [True] * 5, id="high_rate_replenishes"),
|
|
pytest.param(0.0, 5.0, 5, [True] * 5, id="exact_capacity_consumption"),
|
|
],
|
|
)
|
|
def test_consumption_patterns(
|
|
self,
|
|
rate: float,
|
|
capacity: float,
|
|
consume_count: int,
|
|
expected_results: list[bool],
|
|
) -> None:
|
|
limiter = TokenBucketLimiter(rate=rate, capacity=capacity)
|
|
results = [limiter.consume("user1") for _ in range(consume_count)]
|
|
assert results == expected_results
|
|
|
|
@pytest.mark.parametrize(
|
|
"tokens_to_consume,capacity,should_succeed",
|
|
[
|
|
pytest.param(5.0, 10.0, True, id="consume_half_capacity"),
|
|
pytest.param(10.0, 10.0, True, id="consume_full_capacity"),
|
|
pytest.param(11.0, 10.0, False, id="consume_over_capacity"),
|
|
pytest.param(0.0, 10.0, True, id="consume_zero"),
|
|
],
|
|
)
|
|
def test_multi_token_consumption(
|
|
self,
|
|
tokens_to_consume: float,
|
|
capacity: float,
|
|
should_succeed: bool,
|
|
) -> None:
|
|
limiter = TokenBucketLimiter(rate=0.0, capacity=capacity)
|
|
assert limiter.consume("user1", tokens=tokens_to_consume) is should_succeed
|
|
|
|
def test_independent_keys(self) -> None:
|
|
limiter = TokenBucketLimiter(rate=0.0, capacity=1.0)
|
|
|
|
assert limiter.consume("user1") is True
|
|
assert limiter.consume("user1") is False
|
|
assert limiter.consume("user2") is True
|
|
assert limiter.consume("user2") is False
|
|
|
|
def test_clear_resets_all_buckets(self) -> None:
|
|
limiter = TokenBucketLimiter(rate=0.0, capacity=1.0)
|
|
|
|
limiter.consume("user1")
|
|
limiter.consume("user2")
|
|
assert limiter.consume("user1") is False
|
|
assert limiter.consume("user2") is False
|
|
|
|
limiter.clear()
|
|
|
|
assert limiter.consume("user1") is True
|
|
assert limiter.consume("user2") is True
|
|
|
|
def test_tokens_replenish_over_time(self) -> None:
|
|
limiter = TokenBucketLimiter(rate=10.0, capacity=1.0)
|
|
|
|
assert limiter.consume("user1") is True
|
|
assert limiter.consume("user1") is False
|
|
|
|
time.sleep(0.15)
|
|
|
|
assert limiter.consume("user1") is True
|
|
|
|
def test_partial_replenishment(self) -> None:
|
|
limiter = TokenBucketLimiter(rate=10.0, capacity=2.0)
|
|
|
|
assert limiter.consume("user1", tokens=2.0) is True
|
|
assert limiter.consume("user1") is False
|
|
|
|
time.sleep(0.1)
|
|
|
|
assert limiter.consume("user1") is True
|
|
assert limiter.consume("user1") is False
|
|
|
|
|
|
class TestRateLimiter:
|
|
@pytest.fixture
|
|
def fake_redis(self) -> fakeredis.FakeRedis:
|
|
return fakeredis.FakeRedis()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_only_when_no_redis(self) -> None:
|
|
limiter = RateLimiter(
|
|
redis=None,
|
|
burst_limit=2,
|
|
burst_window=60,
|
|
sustained_limit=10,
|
|
sustained_window=3600,
|
|
)
|
|
|
|
allowed, scope = await limiter.check(user_id=1)
|
|
assert allowed is True
|
|
assert scope is None
|
|
|
|
await limiter.check(user_id=1)
|
|
allowed, scope = await limiter.check(user_id=1)
|
|
assert allowed is False
|
|
assert scope == "burst"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_rate_limiting(self, fake_redis: fakeredis.FakeRedis) -> None:
|
|
limiter = RateLimiter(
|
|
redis=fake_redis,
|
|
burst_limit=2,
|
|
burst_window=60,
|
|
sustained_limit=10,
|
|
sustained_window=3600,
|
|
)
|
|
|
|
allowed, _ = await limiter.check(user_id=1)
|
|
assert allowed is True
|
|
|
|
await limiter.check(user_id=1)
|
|
allowed, scope = await limiter.check(user_id=1)
|
|
assert allowed is False
|
|
assert scope == "burst"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_keys_are_set_with_expiry(self, fake_redis: fakeredis.FakeRedis) -> None:
|
|
limiter = RateLimiter(
|
|
redis=fake_redis,
|
|
burst_limit=100,
|
|
burst_window=60,
|
|
sustained_limit=1000,
|
|
sustained_window=3600,
|
|
)
|
|
|
|
await limiter.check(user_id=123)
|
|
|
|
burst_ttl = await fake_redis.ttl("ratelimit:burst:123")
|
|
sustained_ttl = await fake_redis.ttl("ratelimit:sustained:123")
|
|
|
|
assert 0 < burst_ttl <= 60
|
|
assert 0 < sustained_ttl <= 3600
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fallback_to_local_on_redis_error(self) -> None:
|
|
mock_redis = MagicMock()
|
|
mock_redis.incr = AsyncMock(side_effect=Exception("Redis connection failed"))
|
|
|
|
limiter = RateLimiter(
|
|
redis=mock_redis,
|
|
burst_limit=2,
|
|
burst_window=60,
|
|
sustained_limit=10,
|
|
sustained_window=3600,
|
|
)
|
|
|
|
allowed, scope = await limiter.check(user_id=1)
|
|
assert allowed is True
|
|
assert scope is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_independent_users(self, fake_redis: fakeredis.FakeRedis) -> None:
|
|
limiter = RateLimiter(
|
|
redis=fake_redis,
|
|
burst_limit=1,
|
|
burst_window=60,
|
|
sustained_limit=10,
|
|
sustained_window=3600,
|
|
)
|
|
|
|
allowed1, _ = await limiter.check(user_id=1)
|
|
allowed2, _ = await limiter.check(user_id=2)
|
|
|
|
assert allowed1 is True
|
|
assert allowed2 is True
|
|
|
|
denied1, _ = await limiter.check(user_id=1)
|
|
denied2, _ = await limiter.check(user_id=2)
|
|
|
|
assert denied1 is False
|
|
assert denied2 is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sustained_limit_after_burst_ok(self, fake_redis: fakeredis.FakeRedis) -> None:
|
|
limiter = RateLimiter(
|
|
redis=fake_redis,
|
|
burst_limit=100,
|
|
burst_window=60,
|
|
sustained_limit=2,
|
|
sustained_window=3600,
|
|
)
|
|
|
|
await limiter.check(user_id=1)
|
|
await limiter.check(user_id=1)
|
|
allowed, scope = await limiter.check(user_id=1)
|
|
|
|
assert allowed is False
|
|
assert scope == "sustained"
|
|
|
|
|
|
class TestRateLimitResponseHeaders:
|
|
def test_429_includes_retry_after_header_and_structured_error(self, mock_db_pool: MagicMock) -> None:
|
|
class AlwaysDenyThrottle(Throttle):
|
|
scope = "test_throttle"
|
|
|
|
async def allow_request(self, context: ThrottleContext) -> ThrottleResult:
|
|
return ThrottleResult.deny(
|
|
detail="Product rate limit exceeded",
|
|
scope=self.scope,
|
|
retry_after=3600,
|
|
)
|
|
|
|
conn = AsyncMock()
|
|
conn.fetchrow = AsyncMock(
|
|
return_value={
|
|
"id": "key_id",
|
|
"user_id": 1,
|
|
"scopes": ["llm_gateway:read"],
|
|
"current_team_id": 1,
|
|
"distinct_id": "test-distinct-id",
|
|
"is_staff": False,
|
|
}
|
|
)
|
|
mock_db_pool.acquire = AsyncMock(return_value=conn)
|
|
|
|
app = create_test_app(mock_db_pool, throttles=[AlwaysDenyThrottle()])
|
|
|
|
with TestClient(app) as client:
|
|
body = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hi"}]}
|
|
headers = {"Authorization": "Bearer sk-test_key"}
|
|
response = client.post("/v1/chat/completions", json=body, headers=headers)
|
|
|
|
assert response.status_code == 429
|
|
assert response.headers["retry-after"] == "3600"
|
|
# The reason is repeated in the message (SDK error strings often
|
|
# surface only error.message) and the throttle scope rides along as
|
|
# a machine-readable code.
|
|
assert response.json() == {
|
|
"error": {
|
|
"message": "Rate limit exceeded: Product rate limit exceeded",
|
|
"type": "rate_limit_error",
|
|
"reason": "Product rate limit exceeded",
|
|
"code": "test_throttle",
|
|
}
|
|
}
|
|
|
|
|
|
class TestFreeTierModelGateErrorBody:
|
|
def test_gate_403_wire_body_carries_code_and_legacy_shim(
|
|
self, mock_db_pool: MagicMock, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""Pins the gate 403's wire shape end-to-end (raise site + exception-handler
|
|
unwrap): a top-level error envelope, never FastAPI's {"detail": ...} nesting.
|
|
Pre-cutover Insights Desktop clients route this 403 by the "(rate_limit)"
|
|
substring in the message and newer clients read the code — a regression in
|
|
either strands installed builds in fatal-session teardown."""
|
|
from llm_gateway.auth.models import AuthenticatedUser
|
|
from llm_gateway.config import get_settings
|
|
from llm_gateway.dependencies import get_authenticated_user
|
|
from llm_gateway.products.config import INSIGHTS_CODE_US_APP_ID
|
|
|
|
monkeypatch.setenv("LLM_GATEWAY_INSIGHTS_CODE_MODEL_GATE_ENABLED", "true")
|
|
get_settings.cache_clear()
|
|
try:
|
|
app = create_test_app(mock_db_pool)
|
|
# OAuth caller on the Desktop app whose org isn't billed for Desktop usage
|
|
# (the conftest quota resolver reports code_usage_billing_active=False).
|
|
app.dependency_overrides[get_authenticated_user] = lambda: AuthenticatedUser(
|
|
user_id=7,
|
|
team_id=1,
|
|
auth_method="oauth_access_token",
|
|
distinct_id="unbilled-user",
|
|
scopes=["*"],
|
|
application_id=INSIGHTS_CODE_US_APP_ID,
|
|
)
|
|
|
|
with TestClient(app) as client:
|
|
response = client.post(
|
|
"/insights_code/v1/messages",
|
|
json={
|
|
"model": "claude-fable-5",
|
|
"max_tokens": 16,
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 403
|
|
body = response.json()
|
|
assert "detail" not in body
|
|
error = body["error"]
|
|
assert error["type"] == "permission_error"
|
|
assert error["code"] == "model_gate"
|
|
assert "claude-fable-5" in error["message"]
|
|
assert error["message"].endswith("(rate_limit)")
|
|
finally:
|
|
get_settings.cache_clear()
|