mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 16:16:48 +00:00
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:
1 parent
6826313b60
commit
5406550bca
13 files changed
+238
-9
No files matched your search
@@ -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,
|
||||
):
|
||||
|
||||
Reference in new issue
Block a user