feat(quotas): enforce user quotas on chat, agent, scheduled and webhook runs

check_usage now checks the billable user's quota on every request, before
the per-agent 24h limits, which keep applying to traffic through an agent.
Until now a request without an agent key skipped every limit. A refusal is
a 429 with Retry-After and a body naming the budget, usage, limit, the
layer the limit came from and when it resets.

Headless runs check the agent owner's quota before starting. A refused
scheduled run is recorded as budget_exceeded; a refused webhook run returns
a quota_exceeded result instead of raising, so Celery does not retry it.
This commit is contained in:
Alex committed 2026-09-21 12:11:11 +01:00
1 parent 6826313b60
commit 5406550bca
13 files changed
+238 -9

No files matched your search

+25
View File
@@ -126,6 +126,31 @@ class TestExecuteScheduledRunBody:
assert sched["consecutive_failure_count"] == 1
assert "schedule.run.failed" in {e[0] for e in stub_events}
def test_quota_refusal_marks_budget_exceeded(
self, pg_engine, patched_engine, stub_events,
):
from datetime import datetime, timezone
from docsgpt.quotas.service import QuotaExceeded, QuotaExceededError
exceeded = QuotaExceeded(
user_id="u1", bucket="all", budget="tokens", usage=10, limit=10,
source="instance", source_id=None,
resets_at=datetime(2099, 1, 1, tzinfo=timezone.utc),
)
with pg_engine.begin() as conn:
_, run, _ = _make_pending_run(conn)
with patch(
"docsgpt.api.user.scheduler_worker.run_agent_headless",
side_effect=QuotaExceededError(exceeded),
):
result = execute_scheduled_run_body(str(run["id"]), "celery-q")
assert result["status"] == "failed"
with pg_engine.connect() as conn:
row = ScheduleRunsRepository(conn).get_internal(str(run["id"]))
assert row["error_type"] == "budget_exceeded"
assert "Usage quota reached" in row["error"]
def test_autopause_after_threshold(
self, pg_engine, patched_engine, stub_events,
):