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.

docsgpt/agents/headless_runner.py

@@ -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")

docsgpt/api/answer/routes/answer.py

@@ -103,7 +103,9 @@ class AnswerResource(Resource, BaseAnswerResource):
                 )
                 if not processor.decoded_token:
                     return make_response({"error": "Unauthorized"}, 401)
-                if error := self.check_usage(processor.agent_config):
+                if error := self.check_usage(
+                    processor.agent_config, processor.decoded_token
+                ):
                     return error
                 stream = self.complete_stream(
                     question="",
@@ -129,7 +131,9 @@ class AnswerResource(Resource, BaseAnswerResource):
                 if not processor.decoded_token:
                     return make_response({"error": "Unauthorized"}, 401)
 
-                if error := self.check_usage(processor.agent_config):
+                if error := self.check_usage(
+                    processor.agent_config, processor.decoded_token
+                ):
                     return error
 
                 should_persist, visibility = resolve_persistence(

docsgpt/api/answer/routes/base.py

@@ -23,6 +23,8 @@ from docsgpt.core.model_utils import (
 from docsgpt.core.settings import settings
 from docsgpt.error import sanitize_api_error
 from docsgpt.llm.llm_creator import LLMCreator
+from docsgpt.quotas.http import quota_exceeded_response
+from docsgpt.quotas.service import QuotaService
 from docsgpt.storage.db.repositories.agents import AgentsRepository
 from docsgpt.storage.db.repositories.conversations import (
     HeartbeatState,
@@ -112,17 +114,29 @@ class BaseAnswerResource:
             prepared.append(item)
         return prepared
 
-    def check_usage(self, agent_config: Dict) -> Optional[Response]:
-        """Check if there is a usage limit and if it is exceeded
+    def check_usage(
+        self, agent_config: Dict, decoded_token: Optional[Dict] = None
+    ) -> Optional[Response]:
+        """Refuse the request when a usage limit is exhausted.
+
+        The billable user's quota is checked first, for every request; the
+        agent's own 24h token and request limits then apply to traffic that
+        runs through an agent.
 
         Args:
             agent_config: The config dict of agent instance
+            decoded_token: The request's resolved identity; its ``sub`` is the
+                billable user.
 
         Returns:
             None or Response if either of limits exceeded.
 
         """
         api_key = agent_config.get("user_api_key")
+        user_id = (decoded_token or {}).get("sub") or agent_config.get("user_id")
+        exceeded = QuotaService.check(user_id, "agent" if api_key else "direct")
+        if exceeded is not None:
+            return quota_exceeded_response(exceeded)
         if not api_key:
             return None
         with db_readonly() as conn:

docsgpt/api/answer/routes/stream.py

@@ -115,7 +115,9 @@ class StreamResource(Resource, BaseAnswerResource):
                         status=401,
                         mimetype="text/event-stream",
                     )
-                if error := self.check_usage(processor.agent_config):
+                if error := self.check_usage(
+                processor.agent_config, processor.decoded_token
+            ):
                     return error
                 return Response(
                     with_sse_keepalive(
@@ -151,7 +153,9 @@ class StreamResource(Resource, BaseAnswerResource):
                     mimetype="text/event-stream",
                 )
 
-            if error := self.check_usage(processor.agent_config):
+            if error := self.check_usage(
+                processor.agent_config, processor.decoded_token
+            ):
                 return error
             should_persist, visibility = resolve_persistence(
                 visibility_flag=data.get("visibility"),

docsgpt/api/user/scheduler_worker.py

@@ -17,6 +17,7 @@ from sqlalchemy import text as sql_text
 from docsgpt.agents.headless_runner import run_agent_headless
 from docsgpt.core.settings import settings
 from docsgpt.events.publisher import publish_user_event
+from docsgpt.quotas.service import QuotaExceededError
 from docsgpt.storage.db.base_repository import row_to_dict
 from docsgpt.storage.db.engine import get_engine
 from docsgpt.storage.db.repositories.conversations import (
@@ -282,6 +283,11 @@ def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Di
         outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
         error_type = "timeout"
         error_text = "run exceeded soft time limit"
+    except QuotaExceededError as exc:
+        # The owner's usage quota is spent; the run never started.
+        outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
+        error_type = "budget_exceeded"
+        error_text = str(exc)
     except Exception as exc:
         outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
         error_type = "agent_error"

docsgpt/api/v1/routes.py

@@ -338,7 +338,7 @@ def chat_completions():
             )
 
         helper = _V1AnswerHelper()
-        usage_error = helper.check_usage(processor.agent_config)
+        usage_error = helper.check_usage(processor.agent_config, processor.decoded_token)
         if usage_error:
             return usage_error
 

docsgpt/quotas/__init__.py

@@ -8,13 +8,14 @@ with their ``token_usage`` totals over the current ``QUOTA_PERIOD`` window.
 
 from docsgpt.quotas.providers import QuotaDefaultsProvider, register_defaults_provider
 from docsgpt.quotas.resolver import ResolvedLimit, ResolvedLimits, resolve_limits
-from docsgpt.quotas.service import BucketStatus, QuotaExceeded, QuotaService
+from docsgpt.quotas.service import BucketStatus, QuotaExceeded, QuotaExceededError, QuotaService
 from docsgpt.quotas.windows import window_bounds
 
 __all__ = [
     "BucketStatus",
     "QuotaDefaultsProvider",
     "QuotaExceeded",
+    "QuotaExceededError",
     "QuotaService",
     "ResolvedLimit",
     "ResolvedLimits",

docsgpt/quotas/http.py

@@ -0,0 +1,14 @@
+"""HTTP rendering of a quota refusal."""
+
+from __future__ import annotations
+
+from flask import Response, jsonify, make_response
+
+from docsgpt.quotas.service import QuotaExceeded
+
+
+def quota_exceeded_response(exceeded: QuotaExceeded) -> Response:
+    """Return the 429 for an exhausted quota, with ``Retry-After`` set to the reset."""
+    response = make_response(jsonify(exceeded.to_payload()), 429)
+    response.headers["Retry-After"] = str(exceeded.retry_after_seconds)
+    return response

docsgpt/quotas/service.py

@@ -104,6 +104,14 @@ class QuotaExceeded:
             return payload
 
 
+class QuotaExceededError(Exception):
+    """Raised where a refused request has no HTTP response to carry the refusal."""
+
+    def __init__(self, exceeded: QuotaExceeded) -> None:
+        self.exceeded = exceeded
+        super().__init__(exceeded.to_payload()["message"])
+
+
 class QuotaService:
     """Resolve limits and measure usage for the current ``QUOTA_PERIOD`` window."""
 

docsgpt/worker.py

@@ -2122,6 +2122,7 @@ def agent_webhook_worker(self, agent_id, payload):
     try:
         # Shared headless path with the scheduler; approval-gated tools auto-deny.
         from docsgpt.agents.headless_runner import run_agent_headless
+        from docsgpt.quotas.service import QuotaExceededError
 
         outcome = run_agent_headless(
             agent_config,
@@ -2135,6 +2136,12 @@ def agent_webhook_worker(self, agent_id, payload):
             "tool_calls": outcome.get("tool_calls", []),
             "thought": outcome.get("thought", ""),
         }
+    except QuotaExceededError as e:
+        # Returned, not raised: retrying cannot succeed before the quota resets.
+        logging.warning(
+            f"Webhook skipped for agent {agent_id}: {e}", extra={"agent_id": agent_id}
+        )
+        return {"status": "quota_exceeded", "error": str(e)}
     except Exception as e:
         logging.error(f"Error running agent logic: {e}", exc_info=True)
         raise