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.
486 lines
20 KiB
Python
486 lines
20 KiB
Python
from collections.abc import Generator
|
|
from datetime import UTC, datetime, timedelta
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
import structlog
|
|
from fastapi.testclient import TestClient
|
|
|
|
from llm_gateway.api.usage import _to_cost_limit_status
|
|
from llm_gateway.rate_limiting.cost_throttles import (
|
|
CostStatus,
|
|
UserCostBurstThrottle,
|
|
UserCostSustainedThrottle,
|
|
)
|
|
from llm_gateway.services.plan_resolver import BillingPeriod, PlanInfo
|
|
from tests.conftest import create_test_app
|
|
|
|
|
|
class TestToCostLimitStatus:
|
|
NOW = datetime(2026, 5, 1, 12, 0, 0, tzinfo=UTC)
|
|
|
|
@pytest.mark.parametrize(
|
|
"used_usd,limit_usd,expected_percent",
|
|
[
|
|
(0.0, 100.0, 0.0),
|
|
(25.5, 100.0, 25.5),
|
|
(100.0, 100.0, 100.0),
|
|
(150.0, 100.0, 100.0),
|
|
(0.0, 0.0, 100.0),
|
|
(5.0, 0.0, 100.0),
|
|
],
|
|
)
|
|
def test_used_percent(self, used_usd: float, limit_usd: float, expected_percent: float) -> None:
|
|
status = CostStatus(
|
|
used_usd=used_usd, limit_usd=limit_usd, remaining_usd=0.0, resets_in_seconds=60, exceeded=False
|
|
)
|
|
result = _to_cost_limit_status(status, now=self.NOW)
|
|
assert result.used_percent == expected_percent
|
|
|
|
def test_passes_through_exceeded_and_resets(self) -> None:
|
|
status = CostStatus(used_usd=100.0, limit_usd=100.0, remaining_usd=0.0, resets_in_seconds=3600, exceeded=True)
|
|
result = _to_cost_limit_status(status, now=self.NOW)
|
|
assert result.exceeded is True
|
|
assert result.resets_in_seconds == 3600
|
|
|
|
@pytest.mark.parametrize("resets_in_seconds", [0, 60, 3600, 86400])
|
|
def test_reset_at_matches_resets_in_seconds(self, resets_in_seconds: int) -> None:
|
|
status = CostStatus(
|
|
used_usd=10.0,
|
|
limit_usd=100.0,
|
|
remaining_usd=90.0,
|
|
resets_in_seconds=resets_in_seconds,
|
|
exceeded=False,
|
|
)
|
|
result = _to_cost_limit_status(status, now=self.NOW)
|
|
assert result.reset_at == self.NOW + timedelta(seconds=resets_in_seconds)
|
|
|
|
|
|
class TestUsageEndpoint:
|
|
@pytest.fixture
|
|
def authenticated_usage_client(self, mock_db_pool: MagicMock) -> Generator[TestClient]:
|
|
app = create_test_app(mock_db_pool)
|
|
|
|
conn = AsyncMock()
|
|
conn.fetchrow = AsyncMock(
|
|
return_value={
|
|
"id": "key_id",
|
|
"user_id": 42,
|
|
"scopes": ["llm_gateway:read"],
|
|
"current_team_id": 1,
|
|
"distinct_id": "test-distinct-id",
|
|
"is_staff": False,
|
|
}
|
|
)
|
|
mock_db_pool.acquire = AsyncMock(return_value=conn)
|
|
mock_db_pool.release = AsyncMock()
|
|
|
|
with TestClient(app) as c:
|
|
yield c
|
|
|
|
def test_returns_pro_limits_by_default(self, authenticated_usage_client: TestClient) -> None:
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["product"] == "insights_code"
|
|
assert data["user_id"] == 42
|
|
assert data["burst"]["used_percent"] == 0
|
|
assert data["sustained"]["used_percent"] == 0
|
|
assert data["is_rate_limited"] is False
|
|
assert "is_pro" in data
|
|
assert "reset_at" in data["burst"]
|
|
assert "reset_at" in data["sustained"]
|
|
assert "billing_period_end" in data
|
|
|
|
def test_response_has_no_usd_fields(self, authenticated_usage_client: TestClient) -> None:
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
for bucket in (data["burst"], data["sustained"]):
|
|
for key in bucket:
|
|
assert "usd" not in key.lower(), f"Bucket field {key!r} leaks USD"
|
|
|
|
def test_includes_billing_period_end_for_pro_plan(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(
|
|
plan_key="insights-code-200-20260301",
|
|
seat_created_at="2026-01-01T00:00:00+00:00",
|
|
billing_period=BillingPeriod(
|
|
current_period_start="2026-05-01T00:00:00+00:00",
|
|
current_period_end="2026-05-31T00:00:00+00:00",
|
|
interval="month",
|
|
),
|
|
)
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["billing_period_end"] is not None
|
|
assert data["billing_period_end"].startswith("2026-05-31")
|
|
|
|
def test_billing_period_end_normalises_naive_iso_to_utc(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(
|
|
plan_key="insights-code-200-20260301",
|
|
seat_created_at="2026-01-01T00:00:00+00:00",
|
|
billing_period=BillingPeriod(
|
|
current_period_start="2026-05-01T00:00:00",
|
|
current_period_end="2026-05-31T00:00:00",
|
|
interval="month",
|
|
),
|
|
)
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
billing_period_end = response.json()["billing_period_end"]
|
|
assert billing_period_end is not None
|
|
parsed = datetime.fromisoformat(billing_period_end)
|
|
assert parsed.tzinfo is not None
|
|
assert parsed.utcoffset() == timedelta(0)
|
|
|
|
def test_billing_period_end_null_when_unparseable(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(
|
|
plan_key="insights-code-200-20260301",
|
|
seat_created_at="2026-01-01T00:00:00+00:00",
|
|
billing_period=BillingPeriod(
|
|
current_period_start="2026-05-01T00:00:00+00:00",
|
|
current_period_end="not-a-real-date",
|
|
interval="month",
|
|
),
|
|
)
|
|
)
|
|
|
|
with structlog.testing.capture_logs() as logs:
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["billing_period_end"] is None
|
|
assert any(log.get("event") == "usage.billing_period_end_unparseable" for log in logs)
|
|
|
|
def test_billing_period_end_null_for_free_plan(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(plan_key=None, seat_created_at="2026-01-01T00:00:00+00:00")
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["billing_period_end"] is None
|
|
|
|
@pytest.mark.parametrize(
|
|
"plan_key,expected_is_pro",
|
|
[
|
|
("insights-code-200-20260301", True),
|
|
("insights-code-free-20260301", False),
|
|
(None, False),
|
|
],
|
|
)
|
|
def test_is_pro_reflects_plan_key(
|
|
self, authenticated_usage_client: TestClient, plan_key: str | None, expected_is_pro: bool
|
|
) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(plan_key=plan_key, seat_created_at="2026-01-01T00:00:00+00:00")
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["is_pro"] is expected_is_pro
|
|
|
|
def test_returns_free_limits_for_free_plan_with_seat(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(plan_key=None, seat_created_at="2026-01-01T00:00:00+00:00")
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["burst"]["used_percent"] == 0
|
|
assert data["sustained"]["used_percent"] == 0
|
|
|
|
def test_old_free_user_still_gets_limits(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(
|
|
plan_key="insights-code-free-20260301",
|
|
seat_created_at="2025-01-01T00:00:00+00:00",
|
|
)
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["burst"]["used_percent"] == 0
|
|
assert data["burst"]["exceeded"] is False
|
|
assert data["sustained"]["used_percent"] == 0
|
|
assert data["sustained"]["exceeded"] is False
|
|
assert data["is_rate_limited"] is False
|
|
|
|
def test_returns_pro_limits_with_pro_plan(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(plan_key="insights-code-200-20260301", seat_created_at=None)
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["burst"]["used_percent"] == 0
|
|
assert data["sustained"]["used_percent"] == 0
|
|
assert data["is_rate_limited"] is False
|
|
|
|
def test_returns_401_without_auth(self, client: TestClient) -> None:
|
|
response = client.get("/v1/usage/insights_code")
|
|
assert response.status_code == 401
|
|
|
|
def test_reflects_accumulated_cost(
|
|
self, authenticated_usage_client: TestClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
app = authenticated_usage_client.app
|
|
runner = app.state.throttle_runner
|
|
|
|
burst_throttle = next(t for t in runner.throttles if isinstance(t, UserCostBurstThrottle))
|
|
sustained_throttle = next(t for t in runner.throttles if isinstance(t, UserCostSustainedThrottle))
|
|
|
|
monkeypatch.setattr(
|
|
burst_throttle,
|
|
"get_status",
|
|
AsyncMock(
|
|
return_value=CostStatus(
|
|
used_usd=25.5, limit_usd=100.0, remaining_usd=74.5, resets_in_seconds=3600, exceeded=False
|
|
)
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
sustained_throttle,
|
|
"get_status",
|
|
AsyncMock(
|
|
return_value=CostStatus(
|
|
used_usd=25.5, limit_usd=1000.0, remaining_usd=974.5, resets_in_seconds=86400, exceeded=False
|
|
)
|
|
),
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["burst"]["used_percent"] == 25.5
|
|
assert data["burst"]["exceeded"] is False
|
|
|
|
assert data["sustained"]["used_percent"] == 2.5
|
|
assert data["sustained"]["exceeded"] is False
|
|
|
|
def test_shows_rate_limited_when_burst_exceeded(
|
|
self, authenticated_usage_client: TestClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
app = authenticated_usage_client.app
|
|
runner = app.state.throttle_runner
|
|
|
|
burst_throttle = next(t for t in runner.throttles if isinstance(t, UserCostBurstThrottle))
|
|
sustained_throttle = next(t for t in runner.throttles if isinstance(t, UserCostSustainedThrottle))
|
|
|
|
monkeypatch.setattr(
|
|
burst_throttle,
|
|
"get_status",
|
|
AsyncMock(
|
|
return_value=CostStatus(
|
|
used_usd=100.0, limit_usd=100.0, remaining_usd=0, resets_in_seconds=3600, exceeded=True
|
|
)
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
sustained_throttle,
|
|
"get_status",
|
|
AsyncMock(
|
|
return_value=CostStatus(
|
|
used_usd=100.0, limit_usd=1000.0, remaining_usd=900.0, resets_in_seconds=86400, exceeded=False
|
|
)
|
|
),
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["burst"]["exceeded"] is True
|
|
assert data["is_rate_limited"] is True
|
|
|
|
def test_ignores_user_id_query_param(self, authenticated_usage_client: TestClient) -> None:
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code?user_id=99",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["user_id"] == 42
|
|
|
|
def test_credits_reflect_products_own_bucket(self, authenticated_usage_client: TestClient) -> None:
|
|
# insights_code's usage reports the insights_code_credits bucket (under the
|
|
# legacy `ai_credits` response field), resolved against its own resource key.
|
|
from llm_gateway.services.quota_resolver import QuotaResourceStatus
|
|
|
|
app = authenticated_usage_client.app
|
|
resolver_mock = AsyncMock(return_value=QuotaResourceStatus(limited=True))
|
|
app.state.quota_resolver.get_resource_status = resolver_mock
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(plan_key=None, seat_created_at=None, seat_missing=True)
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["ai_credits"] == {"exhausted": True, "used_usd": None, "limit_usd": None}
|
|
assert data["is_rate_limited"] is True
|
|
assert resolver_mock.call_args.args[0] == "insights_code_credits"
|
|
|
|
@pytest.mark.parametrize("plan_key", ["insights-code-200-20260301", "insights-code-free-20260301", None])
|
|
def test_exhausted_bucket_reported_for_every_caller(
|
|
self, authenticated_usage_client: TestClient, plan_key: str | None
|
|
) -> None:
|
|
# Mirrors BillableCreditThrottle: every insights_code generation counts into
|
|
# the org's bucket, so exhaustion blocks (and must be reported to) every
|
|
# caller regardless of seat state — clients gate on this response, and a
|
|
# carve-out here would let them keep sending requests enforcement denies.
|
|
from llm_gateway.services.quota_resolver import QuotaResourceStatus
|
|
|
|
app = authenticated_usage_client.app
|
|
app.state.quota_resolver.get_resource_status = AsyncMock(return_value=QuotaResourceStatus(limited=True))
|
|
app.state.plan_resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(plan_key=plan_key, seat_created_at="2026-01-01T00:00:00+00:00")
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["ai_credits"] == {"exhausted": True, "used_usd": None, "limit_usd": None}
|
|
assert data["is_rate_limited"] is True
|
|
|
|
def test_ai_credits_reflects_resolver_for_billable_product(self, authenticated_usage_client: TestClient) -> None:
|
|
from llm_gateway.services.quota_resolver import QuotaResourceStatus
|
|
|
|
app = authenticated_usage_client.app
|
|
resolver_mock = AsyncMock(return_value=QuotaResourceStatus(limited=True))
|
|
app.state.quota_resolver.get_resource_status = resolver_mock
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/slack_app",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["ai_credits"] == {"exhausted": True, "used_usd": None, "limit_usd": None}
|
|
assert data["is_rate_limited"] is True
|
|
assert resolver_mock.call_args.args[0] == "ai_credits"
|
|
|
|
def test_ai_credits_carries_org_spend_numbers(self, authenticated_usage_client: TestClient) -> None:
|
|
"""Insights Desktop renders "used $X of $Y" (titlebar, plans page) off these
|
|
numbers; None must stay None so clients read unknown, not $0."""
|
|
from llm_gateway.services.quota_resolver import QuotaResourceStatus
|
|
|
|
app = authenticated_usage_client.app
|
|
app.state.quota_resolver.get_resource_status = AsyncMock(
|
|
return_value=QuotaResourceStatus(limited=False, used_usd=12.4, limit_usd=50.0)
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["ai_credits"] == {"exhausted": False, "used_usd": 12.4, "limit_usd": 50.0}
|
|
|
|
@pytest.mark.parametrize("billing_active", [True, False])
|
|
def test_code_usage_subscribed_reflects_billing_bit(
|
|
self, authenticated_usage_client: TestClient, billing_active: bool
|
|
) -> None:
|
|
"""Clients pick billing copy and hide the free-tier meter off this
|
|
field, so it must carry the same bit enforcement reads."""
|
|
from llm_gateway.services.quota_resolver import QuotaResourceStatus
|
|
|
|
app = authenticated_usage_client.app
|
|
app.state.quota_resolver.get_resource_status = AsyncMock(
|
|
return_value=QuotaResourceStatus(limited=False, code_usage_billing_active=billing_active)
|
|
)
|
|
|
|
response = authenticated_usage_client.get(
|
|
"/v1/usage/insights_code",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["code_usage_subscribed"] is billing_active
|
|
|
|
def test_invalidate_plan_cache_calls_resolver(self, authenticated_usage_client: TestClient) -> None:
|
|
app = authenticated_usage_client.app
|
|
app.state.plan_resolver.invalidate = AsyncMock()
|
|
|
|
response = authenticated_usage_client.post(
|
|
"/v1/usage/insights_code/invalidate-plan-cache",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json() == {"ok": True}
|
|
app.state.plan_resolver.invalidate.assert_called_once_with(42)
|
|
|
|
def test_invalidate_plan_cache_404_for_other_product(self, authenticated_usage_client: TestClient) -> None:
|
|
response = authenticated_usage_client.post(
|
|
"/v1/usage/wizard/invalidate-plan-cache",
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
assert response.status_code == 404
|
|
|
|
def test_invalidate_plan_cache_401_without_auth(self, client: TestClient) -> None:
|
|
response = client.post("/v1/usage/insights_code/invalidate-plan-cache")
|
|
assert response.status_code == 401
|