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

+9 -1
View File
@@ -15,6 +15,7 @@ from docsgpt.api.answer.services.prompt_renderer import (
)
from docsgpt.api.answer.services.stream_processor import get_prompt
from docsgpt.core.settings import settings
from docsgpt.quotas.service import QuotaExceededError, QuotaService
from docsgpt.retriever.retriever_creator import RetrieverCreator
from docsgpt.storage.db.repositories.sources import SourcesRepository
from docsgpt.storage.db.session import db_readonly
@@ -69,7 +70,11 @@ def run_agent_headless(
chat_history: Optional[List[Dict[str, Any]]] = None,
conversation_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Run an agent with no live client; returns a structured outcome dict."""
"""Run an agent with no live client; returns a structured outcome dict.
Raises:
QuotaExceededError: If the agent owner's usage quota is exhausted.
"""
from docsgpt.core.model_utils import (
get_api_key_for_provider,
get_default_model_id,
@@ -82,6 +87,9 @@ def run_agent_headless(
if not owner:
raise ValueError("Agent config is missing user_id; cannot run headless.")
decoded_token = {"sub": owner}
exceeded = QuotaService.check(owner, "agent" if agent_config.get("key") else "direct")
if exceeded is not None:
raise QuotaExceededError(exceeded)
retriever_kind = agent_config.get("retriever", "classic")
source_id = agent_config.get("source_id") or agent_config.get("source")