From b77561288d27f02ac96bec284bc268f405a0b6e6 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 11:38:06 +0100 Subject: [PATCH 01/14] feat(pricing): per-million model rates and a cost module Rename the unused *_cost_per_token capability fields to USD per 1M tokens, add prompt-cache read/write rates, and ship list prices for the hosted catalogs. The old per-token keys still load, scaled, with a warning. docsgpt/pricing.py turns a call's token bins into a USD cost. Models with no declared rate cost $0 unless QUOTA_UNPRICED_RATE_PER_MILLION is set. --- docs/content/Deploying/Settings-Reference.mdx | 17 +++ docsgpt/core/model_settings.py | 11 +- docsgpt/core/model_yaml.py | 34 ++++- docsgpt/core/models/README.md | 6 +- docsgpt/core/models/anthropic.yaml | 6 + docsgpt/core/models/deepseek.yaml | 4 + docsgpt/core/models/google.yaml | 7 + docsgpt/core/models/groq.yaml | 6 + docsgpt/core/models/novita.yaml | 6 + docsgpt/core/models/openai.yaml | 7 + docsgpt/core/models/openrouter.yaml | 6 + docsgpt/core/settings/__init__.py | 2 + docsgpt/core/settings/quotas.py | 36 +++++ docsgpt/pricing.py | 103 ++++++++++++++ tests/core/test_model_settings.py | 6 +- tests/test_pricing.py | 130 ++++++++++++++++++ 16 files changed, 375 insertions(+), 12 deletions(-) create mode 100644 docsgpt/core/settings/quotas.py create mode 100644 docsgpt/pricing.py create mode 100644 tests/test_pricing.py diff --git a/docs/content/Deploying/Settings-Reference.mdx b/docs/content/Deploying/Settings-Reference.mdx index 33ea3b53..d9f59ec7 100644 --- a/docs/content/Deploying/Settings-Reference.mdx +++ b/docs/content/Deploying/Settings-Reference.mdx @@ -1451,6 +1451,23 @@ Type `int`, default `30`, must be `>= 1`. Days guardrail events are kept before the cleanup task removes them. +## Quotas + +Quota window and the treatment of unpriced models. + +### `QUOTA_PERIOD` + +Type `"day" | "week" | "month"`, default `month`. + +Window every usage quota is measured over. Windows are calendar-aligned in UTC: a day starts at 00:00, a week on Monday, a month on the 1st. + +### `QUOTA_UNPRICED_RATE_PER_MILLION` + +Type `list[float]`, default unset. + +Fallback `[input, output]` USD rates per 1M tokens for models that declare no price, e.g. `[0.5, 1.5]`. Unset, such calls are recorded at $0 and only count toward token quotas. + + ## Scheduler Cadence, quotas and timeouts of scheduled runs. diff --git a/docsgpt/core/model_settings.py b/docsgpt/core/model_settings.py index d5e8a5c7..6a2a002f 100644 --- a/docsgpt/core/model_settings.py +++ b/docsgpt/core/model_settings.py @@ -32,8 +32,15 @@ class ModelCapabilities: supports_streaming: bool = True supported_attachment_types: List[str] = field(default_factory=list) context_window: int = 128000 - input_cost_per_token: Optional[float] = None - output_cost_per_token: Optional[float] = None + # USD per 1M tokens; consumed by ``docsgpt/pricing.py``. ``None`` means + # "not declared": the call is recorded at $0 unless + # ``QUOTA_UNPRICED_RATE_PER_MILLION`` is set. + input_cost_per_million: Optional[float] = None + output_cost_per_million: Optional[float] = None + # Rates for the prompt-cache sub-bins of the prompt total. ``None`` bills + # those tokens at ``input_cost_per_million``. + cached_input_cost_per_million: Optional[float] = None + cache_write_cost_per_million: Optional[float] = None # OpenAI reasoning-model effort hint (none/minimal/low/medium/high/xhigh; # the accepted subset is model-dependent). Consumed by OpenAILLM — sent # top-level on Chat Completions and nested under ``reasoning`` on the diff --git a/docsgpt/core/model_yaml.py b/docsgpt/core/model_yaml.py index 72f2d1ef..564c362e 100644 --- a/docsgpt/core/model_yaml.py +++ b/docsgpt/core/model_yaml.py @@ -18,7 +18,7 @@ from pathlib import Path from typing import Dict, List, Optional, Sequence import yaml -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from docsgpt.core.model_settings import ( AvailableModel, @@ -65,11 +65,33 @@ class _CapabilityFields(BaseModel): supports_streaming: Optional[bool] = None attachments: Optional[List[str]] = None context_window: Optional[int] = None - input_cost_per_token: Optional[float] = None - output_cost_per_token: Optional[float] = None + input_cost_per_million: Optional[float] = Field(default=None, ge=0) + output_cost_per_million: Optional[float] = Field(default=None, ge=0) + cached_input_cost_per_million: Optional[float] = Field(default=None, ge=0) + cache_write_cost_per_million: Optional[float] = Field(default=None, ge=0) reasoning_effort: Optional[str] = None api_flavor: Optional[str] = None + @model_validator(mode="before") + @classmethod + def _per_token_alias(cls, data): + """Accept the deprecated ``*_cost_per_token`` keys, scaled to per-1M.""" + if not isinstance(data, dict): + return data + data = dict(data) + for side in ("input", "output"): + old, new = f"{side}_cost_per_token", f"{side}_cost_per_million" + if old not in data: + continue + value = data.pop(old) + if new in data: + raise ValueError(f"set only one of {old} and {new}") + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{old} must be a number") + logger.warning("%s is deprecated; use %s (USD per 1M tokens)", old, new) + data[new] = value * 1_000_000 + return data + @field_validator("reasoning_effort") @classmethod def _valid_reasoning_effort(cls, v: Optional[str]) -> Optional[str]: @@ -237,8 +259,10 @@ def _build_model( supports_streaming=pick("supports_streaming", True), supported_attachment_types=expanded, context_window=pick("context_window", 128000), - input_cost_per_token=pick("input_cost_per_token", None), - output_cost_per_token=pick("output_cost_per_token", None), + input_cost_per_million=pick("input_cost_per_million", None), + output_cost_per_million=pick("output_cost_per_million", None), + cached_input_cost_per_million=pick("cached_input_cost_per_million", None), + cache_write_cost_per_million=pick("cache_write_cost_per_million", None), reasoning_effort=pick("reasoning_effort", None), api_flavor=pick("api_flavor", "chat_completions"), ) diff --git a/docsgpt/core/models/README.md b/docsgpt/core/models/README.md index d67f20fe..5e15a382 100644 --- a/docsgpt/core/models/README.md +++ b/docsgpt/core/models/README.md @@ -108,8 +108,10 @@ defaults: # optional, applied to every model below supports_streaming: bool # default true attachments: [, ...] # default [] context_window: int # default 128000 - input_cost_per_token: float # default null - output_cost_per_token: float # default null + input_cost_per_million: float # USD per 1M prompt tokens; default null (unpriced) + output_cost_per_million: float # USD per 1M generated tokens; default null + cached_input_cost_per_million: float # prompt-cache reads; default: the input rate + cache_write_cost_per_million: float # prompt-cache writes; default: the input rate reasoning_effort: # default null; none|minimal|low|medium|high|xhigh (subset is model-dependent) api_flavor: # chat_completions (default) or responses diff --git a/docsgpt/core/models/anthropic.yaml b/docsgpt/core/models/anthropic.yaml index 4e253396..784ab8e3 100644 --- a/docsgpt/core/models/anthropic.yaml +++ b/docsgpt/core/models/anthropic.yaml @@ -10,14 +10,20 @@ models: description: Most capable Claude model for complex reasoning and agentic coding context_window: 1000000 supports_structured_output: true + input_cost_per_million: 5.0 + output_cost_per_million: 25.0 - id: claude-sonnet-4-6 display_name: Claude Sonnet 4.6 description: Best balance of speed and intelligence with extended thinking context_window: 1000000 supports_structured_output: true + input_cost_per_million: 3.0 + output_cost_per_million: 15.0 - id: claude-haiku-4-5 display_name: Claude Haiku 4.5 description: Fastest Claude model with near-frontier intelligence supports_structured_output: true + input_cost_per_million: 1.0 + output_cost_per_million: 5.0 diff --git a/docsgpt/core/models/deepseek.yaml b/docsgpt/core/models/deepseek.yaml index 017c090a..08271e09 100644 --- a/docsgpt/core/models/deepseek.yaml +++ b/docsgpt/core/models/deepseek.yaml @@ -12,7 +12,11 @@ models: - id: deepseek-v4-flash display_name: DeepSeek V4 Flash description: Cost-efficient 1M-context model with hybrid thinking / non-thinking modes, tool calling and FIM completion + input_cost_per_million: 0.14 + output_cost_per_million: 0.28 - id: deepseek-v4-pro display_name: DeepSeek V4 Pro description: Frontier 1M-context model with hybrid thinking / non-thinking modes for advanced reasoning and agentic coding + input_cost_per_million: 0.435 + output_cost_per_million: 0.87 diff --git a/docsgpt/core/models/google.yaml b/docsgpt/core/models/google.yaml index a4102d77..3fe7dbc5 100644 --- a/docsgpt/core/models/google.yaml +++ b/docsgpt/core/models/google.yaml @@ -9,9 +9,16 @@ models: - id: gemini-3.1-pro-preview display_name: Gemini 3.1 Pro (preview) description: Most capable Gemini 3 model with advanced reasoning and agentic coding (preview) + # Priced at the >200k-token tier; long prompts are common with attachments. + input_cost_per_million: 4.0 + output_cost_per_million: 18.0 - id: gemini-3.5-flash display_name: Gemini 3.5 Flash description: Frontier-class Flash for sustained performance on agentic and coding tasks + input_cost_per_million: 1.5 + output_cost_per_million: 9 - id: gemini-3.1-flash-lite display_name: Gemini 3.1 Flash-Lite description: Cost-efficient frontier-class multimodal model for high-throughput workloads + input_cost_per_million: 0.25 + output_cost_per_million: 1.5 diff --git a/docsgpt/core/models/groq.yaml b/docsgpt/core/models/groq.yaml index 555951ec..c6e28d7a 100644 --- a/docsgpt/core/models/groq.yaml +++ b/docsgpt/core/models/groq.yaml @@ -8,9 +8,15 @@ models: display_name: GPT-OSS 120B description: OpenAI's open-weight 120B flagship served on Groq's LPU hardware; strong general reasoning with strict structured output support supports_structured_output: true + input_cost_per_million: 0.15 + output_cost_per_million: 0.6 - id: llama-3.3-70b-versatile display_name: Llama 3.3 70B Versatile description: Meta's Llama 3.3 70B for general-purpose chat with parallel tool use + input_cost_per_million: 0.59 + output_cost_per_million: 0.79 - id: llama-3.1-8b-instant display_name: Llama 3.1 8B Instant description: Small, very low-latency Llama model (~560 tok/s) with parallel tool use + input_cost_per_million: 0.05 + output_cost_per_million: 0.08 diff --git a/docsgpt/core/models/novita.yaml b/docsgpt/core/models/novita.yaml index 3fa2e89f..08a2cbd8 100644 --- a/docsgpt/core/models/novita.yaml +++ b/docsgpt/core/models/novita.yaml @@ -8,14 +8,20 @@ models: display_name: DeepSeek V4 Pro description: 1.6T MoE (49B active) with 1M context, hybrid CSA/HCA attention, top-tier reasoning and agentic coding context_window: 1048576 + input_cost_per_million: 1.6 + output_cost_per_million: 3.2 - id: moonshotai/kimi-k2.6 display_name: Kimi K2.6 description: 1T-parameter open-weight MoE with native vision/video, multi-step tool calling, and agentic long-horizon execution attachments: [image] context_window: 262144 + input_cost_per_million: 0.8 + output_cost_per_million: 3.4 - id: zai-org/glm-5 display_name: GLM-5 description: Z.AI 754B-parameter MoE with strong general reasoning, function calling, and structured output context_window: 202800 + input_cost_per_million: 1.0 + output_cost_per_million: 3.2 diff --git a/docsgpt/core/models/openai.yaml b/docsgpt/core/models/openai.yaml index e0f209c5..598e7b3c 100644 --- a/docsgpt/core/models/openai.yaml +++ b/docsgpt/core/models/openai.yaml @@ -12,9 +12,16 @@ models: context_window: 1050000 api_flavor: responses reasoning_effort: medium + input_cost_per_million: 5.0 + output_cost_per_million: 30.0 + cached_input_cost_per_million: 0.5 - id: gpt-5.4-mini display_name: GPT-5.4 Mini description: Cost-efficient GPT-5.4-class model for high-volume coding, computer use, and subagent workloads + input_cost_per_million: 0.75 + output_cost_per_million: 4.5 - id: gpt-5.4-nano display_name: GPT-5.4 Nano description: Cheapest GPT-5.4-class model, optimized for simple high-volume tasks where speed and cost matter most + input_cost_per_million: 0.2 + output_cost_per_million: 1.25 diff --git a/docsgpt/core/models/openrouter.yaml b/docsgpt/core/models/openrouter.yaml index 0c28dd30..6ac373ee 100644 --- a/docsgpt/core/models/openrouter.yaml +++ b/docsgpt/core/models/openrouter.yaml @@ -10,6 +10,8 @@ models: description: Free-tier 480B MoE coder model with strong agentic tool use; rate-limited context_window: 262000 attachments: [] + input_cost_per_million: 0.0 + output_cost_per_million: 0.0 - id: deepseek/deepseek-v3.2 display_name: DeepSeek V3.2 @@ -17,9 +19,13 @@ models: context_window: 131072 attachments: [] supports_structured_output: true + input_cost_per_million: 0.23 + output_cost_per_million: 0.34 - id: anthropic/claude-sonnet-4.6 display_name: Claude Sonnet 4.6 (via OpenRouter) description: Frontier Sonnet-class model with 1M context, vision, and extended thinking context_window: 1000000 supports_structured_output: true + input_cost_per_million: 3.0 + output_cost_per_million: 15.0 diff --git a/docsgpt/core/settings/__init__.py b/docsgpt/core/settings/__init__.py index 2af13cea..817281c5 100644 --- a/docsgpt/core/settings/__init__.py +++ b/docsgpt/core/settings/__init__.py @@ -28,6 +28,7 @@ from docsgpt.core.settings.guardrails import GuardrailSettings from docsgpt.core.settings.ingestion import IngestionSettings from docsgpt.core.settings.llm import LLMSettings from docsgpt.core.settings.ocr import OCRSettings +from docsgpt.core.settings.quotas import QuotaSettings from docsgpt.core.settings.retrieval import RetrievalSettings from docsgpt.core.settings.sandbox import SandboxSettings from docsgpt.core.settings.scheduler import SchedulerSettings @@ -54,6 +55,7 @@ SETTINGS_GROUPS: tuple[tuple[str, type[SettingsGroup]], ...] = ( ("Events and devices", EventsSettings), ("Agents", AgentSettings), ("Guardrails", GuardrailSettings), + ("Quotas", QuotaSettings), ("Scheduler", SchedulerSettings), ("Sandbox", SandboxSettings), ("Speech", SpeechSettings), diff --git a/docsgpt/core/settings/quotas.py b/docsgpt/core/settings/quotas.py new file mode 100644 index 00000000..d5d5bb34 --- /dev/null +++ b/docsgpt/core/settings/quotas.py @@ -0,0 +1,36 @@ +"""Admin-set usage quotas and the pricing that feeds their cost budgets.""" + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import Field, field_validator + +from docsgpt.core.settings._shared import SettingsGroup + + +class QuotaSettings(SettingsGroup): + """Quota window and the treatment of unpriced models.""" + + QUOTA_PERIOD: Literal["day", "week", "month"] = Field( + default="month", + description=( + "Window every usage quota is measured over. Windows are calendar-aligned in UTC: " + "a day starts at 00:00, a week on Monday, a month on the 1st." + ), + ) + QUOTA_UNPRICED_RATE_PER_MILLION: Optional[list[float]] = Field( + default=None, + description=( + "Fallback `[input, output]` USD rates per 1M tokens for models that declare no price, " + "e.g. `[0.5, 1.5]`. Unset, such calls are recorded at $0 and only count toward token quotas." + ), + ) + @field_validator("QUOTA_UNPRICED_RATE_PER_MILLION") + @classmethod + def _two_non_negative_rates(cls, v: Optional[list[float]]) -> Optional[list[float]]: + if v is None: + return None + if len(v) != 2 or any(rate < 0 for rate in v): + raise ValueError("QUOTA_UNPRICED_RATE_PER_MILLION must be two non-negative numbers") + return v diff --git a/docsgpt/pricing.py b/docsgpt/pricing.py new file mode 100644 index 00000000..a05b8c1b --- /dev/null +++ b/docsgpt/pricing.py @@ -0,0 +1,103 @@ +"""USD cost of LLM calls, from the per-model rates in the model catalogs.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +from docsgpt.core.settings import settings + + +@dataclass(frozen=True) +class ModelRates: + """USD-per-1M rates for one model; ``None`` cache rates bill at the prompt rate.""" + + prompt: float + generated: float + cached_input: Optional[float] = None + cache_write: Optional[float] = None + + +def _unpriced_rates() -> Optional[ModelRates]: + """Return the operator's fallback rates for undeclared models, if configured.""" + fallback = settings.QUOTA_UNPRICED_RATE_PER_MILLION + if not fallback: + return None + return ModelRates(prompt=float(fallback[0]), generated=float(fallback[1])) + + +def resolve_model_rates(model: Optional[str]) -> Optional[ModelRates]: + """Return the rates for a registry model id. + + Args: + model: Canonical registry id (catalog id, or the UUID of a BYOM record). + + Returns: + The declared rates, the ``QUOTA_UNPRICED_RATE_PER_MILLION`` fallback when the + model declares none, or ``None`` when there is no fallback either. + """ + # Imported lazily: the registry pulls in the provider plugins, whose LLM + # classes import ``docsgpt.usage`` and, through it, this module. + from docsgpt.core.model_registry import ModelRegistry + + entry = ModelRegistry.get_instance().models.get(str(model)) if model else None + if entry is None: + return _unpriced_rates() + caps = entry.capabilities + if caps.input_cost_per_million is None or caps.output_cost_per_million is None: + return _unpriced_rates() + cached = caps.cached_input_cost_per_million + written = caps.cache_write_cost_per_million + return ModelRates( + prompt=float(caps.input_cost_per_million), + generated=float(caps.output_cost_per_million), + cached_input=float(cached) if cached is not None else None, + cache_write=float(written) if written is not None else None, + ) + + +def is_priced(model: Optional[str]) -> bool: + """Return whether calls to ``model`` are recorded with a cost.""" + return resolve_model_rates(model) is not None + + +def cost_from_rates( + rates: ModelRates, + prompt_tokens: int, + generated_tokens: int, + cached_tokens: Optional[int] = 0, + cache_write_tokens: Optional[int] = 0, +) -> float: + """Return the USD cost of one call at ``rates``. + + ``prompt_tokens`` is the provider's billing total; ``cached_tokens`` and + ``cache_write_tokens`` are the parts of it read from or written to the prompt + cache. The sub-bins are clamped to the prompt total, so a malformed report can + never price a call below "everything cached". + """ + prompt_total = max(int(prompt_tokens or 0), 0) + cached = min(max(int(cached_tokens or 0), 0), prompt_total) + written = min(max(int(cache_write_tokens or 0), 0), prompt_total - cached) + regular = prompt_total - cached - written + cached_rate = rates.cached_input if rates.cached_input is not None else rates.prompt + write_rate = rates.cache_write if rates.cache_write is not None else rates.prompt + return ( + regular * rates.prompt + + cached * cached_rate + + written * write_rate + + max(int(generated_tokens or 0), 0) * rates.generated + ) / 1_000_000.0 + + +def compute_cost_usd( + model: Optional[str], + prompt_tokens: int, + generated_tokens: int, + cached_tokens: Optional[int] = 0, + cache_write_tokens: Optional[int] = 0, +) -> float: + """Return the USD cost of one call to ``model``; ``0.0`` when it has no rates.""" + rates = resolve_model_rates(model) + if rates is None: + return 0.0 + return cost_from_rates(rates, prompt_tokens, generated_tokens, cached_tokens, cache_write_tokens) diff --git a/tests/core/test_model_settings.py b/tests/core/test_model_settings.py index 266c9426..288ac6c9 100644 --- a/tests/core/test_model_settings.py +++ b/tests/core/test_model_settings.py @@ -46,8 +46,8 @@ class TestModelCapabilities: assert caps.supports_streaming is True assert caps.supported_attachment_types == [] assert caps.context_window == 128000 - assert caps.input_cost_per_token is None - assert caps.output_cost_per_token is None + assert caps.input_cost_per_million is None + assert caps.output_cost_per_million is None @pytest.mark.unit def test_custom_values(self): @@ -55,7 +55,7 @@ class TestModelCapabilities: supports_tools=True, supports_structured_output=True, context_window=32000, - input_cost_per_token=0.001, + input_cost_per_million=1.0, ) assert caps.supports_tools is True assert caps.context_window == 32000 diff --git a/tests/test_pricing.py b/tests/test_pricing.py new file mode 100644 index 00000000..81aa324f --- /dev/null +++ b/tests/test_pricing.py @@ -0,0 +1,130 @@ +"""Tests for docsgpt/pricing.py and the per-million catalog fields.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from docsgpt import pricing +from docsgpt.core.model_settings import ModelCapabilities +from docsgpt.core.model_yaml import ( + BUILTIN_MODELS_DIR, + ModelYAMLError, + load_model_yamls, +) +from docsgpt.pricing import ModelRates, compute_cost_usd, cost_from_rates + + +def _registry(**models): + entries = {k: SimpleNamespace(capabilities=v) for k, v in models.items()} + return SimpleNamespace(models=entries) + + +@pytest.fixture +def priced_registry(): + caps = ModelCapabilities( + input_cost_per_million=2.0, + output_cost_per_million=10.0, + cached_input_cost_per_million=0.2, + cache_write_cost_per_million=2.5, + ) + bare = ModelCapabilities() + with patch( + "docsgpt.core.model_registry.ModelRegistry.get_instance", + return_value=_registry(priced=caps, bare=bare), + ): + yield + + +@pytest.mark.unit +class TestCostFromRates: + def test_prompt_and_generated(self): + rates = ModelRates(prompt=2.0, generated=10.0) + assert cost_from_rates(rates, 1_000_000, 500_000) == pytest.approx(7.0) + + def test_cache_bins_use_their_rates(self): + rates = ModelRates(prompt=2.0, generated=10.0, cached_input=0.2, cache_write=2.5) + cost = cost_from_rates(rates, 1000, 0, cached_tokens=600, cache_write_tokens=100) + assert cost == pytest.approx((300 * 2.0 + 600 * 0.2 + 100 * 2.5) / 1e6) + + def test_missing_cache_rates_bill_at_prompt_rate(self): + rates = ModelRates(prompt=2.0, generated=10.0) + assert cost_from_rates(rates, 1000, 0, cached_tokens=900) == pytest.approx(1000 * 2.0 / 1e6) + + def test_cache_bins_clamped_to_prompt_total(self): + rates = ModelRates(prompt=2.0, generated=0.0, cached_input=0.0, cache_write=0.0) + assert cost_from_rates(rates, 100, 0, cached_tokens=5000, cache_write_tokens=5000) == 0.0 + assert cost_from_rates(rates, 100, 0, cached_tokens=-5) == pytest.approx(100 * 2.0 / 1e6) + + def test_none_and_negative_counts(self): + rates = ModelRates(prompt=2.0, generated=10.0) + assert cost_from_rates(rates, None, -3, None, None) == 0.0 + + +@pytest.mark.unit +class TestComputeCost: + def test_priced_model(self, priced_registry): + assert compute_cost_usd("priced", 1_000_000, 0) == pytest.approx(2.0) + + @pytest.mark.parametrize("model", ["bare", "unknown", None]) + def test_unpriced_model_is_free_without_fallback(self, priced_registry, model): + with patch.object(pricing.settings, "QUOTA_UNPRICED_RATE_PER_MILLION", None): + assert compute_cost_usd(model, 1_000_000, 1_000_000) == 0.0 + assert pricing.is_priced(model) is False + + def test_unpriced_model_uses_fallback(self, priced_registry): + with patch.object(pricing.settings, "QUOTA_UNPRICED_RATE_PER_MILLION", [0.5, 1.5]): + assert compute_cost_usd("bare", 1_000_000, 1_000_000) == pytest.approx(2.0) + assert pricing.is_priced("bare") is True + + +@pytest.mark.unit +class TestCatalogFields: + def _load(self, tmp_path, body): + (tmp_path / "p.yaml").write_text(body) + return load_model_yamls([tmp_path])[0].models[0].capabilities + + def test_per_million_fields(self, tmp_path): + caps = self._load( + tmp_path, + "provider: openai\nmodels:\n - id: m\n input_cost_per_million: 3\n" + " output_cost_per_million: 15\n cached_input_cost_per_million: 0.3\n", + ) + assert (caps.input_cost_per_million, caps.output_cost_per_million) == (3, 15) + assert caps.cached_input_cost_per_million == 0.3 + assert caps.cache_write_cost_per_million is None + + def test_per_token_alias_is_scaled(self, tmp_path): + caps = self._load( + tmp_path, + "provider: openai\ndefaults:\n input_cost_per_token: 0.000003\n" + "models:\n - id: m\n output_cost_per_token: 0.000015\n", + ) + assert caps.input_cost_per_million == pytest.approx(3.0) + assert caps.output_cost_per_million == pytest.approx(15.0) + + def test_both_spellings_rejected(self, tmp_path): + with pytest.raises(ModelYAMLError): + self._load( + tmp_path, + "provider: openai\nmodels:\n - id: m\n input_cost_per_token: 0.1\n" + " input_cost_per_million: 1\n", + ) + + def test_negative_rate_rejected(self, tmp_path): + with pytest.raises(ModelYAMLError): + self._load(tmp_path, "provider: openai\nmodels:\n - id: m\n input_cost_per_million: -1\n") + + def test_hosted_builtin_models_are_priced(self): + hosted = {"anthropic", "deepseek", "google", "groq", "novita", "openai", "openrouter"} + catalogs = [ + c for c in load_model_yamls([BUILTIN_MODELS_DIR]) if c.source_path.stem in hosted + ] + assert {c.source_path.stem for c in catalogs} == hosted + for catalog in catalogs: + for model in catalog.models: + caps = model.capabilities + assert caps.input_cost_per_million is not None, model.id + assert caps.output_cost_per_million is not None, model.id From 43fad2a8653e90e98a39cef1b22cd74ce30d5636 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 11:40:47 +0100 Subject: [PATCH 02/14] feat(quotas): quota_policies table and a per-call cost on token_usage Migration 0033 adds quota_policies (instance default, team per-member allowance, user override; a token budget and a USD budget per row) and token_usage.cost. The column is added IF NOT EXISTS so a database that already carries it upgrades cleanly. Every usage row now records the call's USD cost from the model catalog; bring-your-own models are recorded at $0. --- docsgpt/alembic/versions/0033_quotas.py | 102 ++++++++++++ docsgpt/llm/llm_creator.py | 3 + docsgpt/storage/db/models.py | 43 ++++++ .../storage/db/repositories/token_usage.py | 6 +- docsgpt/usage.py | 25 ++- tests/llm/test_fallback.py | 32 ++++ .../db/repositories/test_token_usage.py | 9 ++ tests/storage/db/test_migration_0033.py | 145 ++++++++++++++++++ tests/test_usage.py | 53 +++++++ 9 files changed, 415 insertions(+), 3 deletions(-) create mode 100644 docsgpt/alembic/versions/0033_quotas.py create mode 100644 tests/storage/db/test_migration_0033.py diff --git a/docsgpt/alembic/versions/0033_quotas.py b/docsgpt/alembic/versions/0033_quotas.py new file mode 100644 index 00000000..165ae021 --- /dev/null +++ b/docsgpt/alembic/versions/0033_quotas.py @@ -0,0 +1,102 @@ +"""0033 quotas — admin-set usage limits and a per-call cost. + +``quota_policies`` holds the limits an instance admin sets at three layers: +the instance default (``subject_id`` NULL), a team's per-member allowance +(``subject_id`` = ``teams.id``) and a single user's override (``subject_id`` = +the auth ``sub``). Each row carries a token budget and a USD budget; per +budget a row either sets a limit (0 blocks), marks it unlimited, or leaves +both empty to defer to the next layer. ``bucket`` narrows a row to chat without +an agent (``direct``) or traffic through an agent (``agent``); ``all`` covers both. + +``subject_id`` is polymorphic, so there is no FK: an AFTER DELETE trigger on +``teams`` scrubs a deleted team's rows, and user rows follow the ``user_roles`` +convention of never blocking user deletion. + +``token_usage.cost`` is the USD cost of the call at write time (see +``docsgpt/pricing.py``); 0 for unpriced and bring-your-own models. +Idempotent both ways. + +Revision ID: 0033_quotas +Revises: 0032_personal_access_tokens +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0033_quotas" +down_revision: Union[str, None] = "0032_personal_access_tokens" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + "ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cost NUMERIC(12,8) NOT NULL DEFAULT 0;" + ) + op.execute( + """ + CREATE TABLE IF NOT EXISTS quota_policies ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + scope TEXT NOT NULL CHECK (scope IN ('instance', 'team', 'user')), + subject_id TEXT, + bucket TEXT NOT NULL DEFAULT 'all' + CHECK (bucket IN ('all', 'direct', 'agent')), + token_limit BIGINT CHECK (token_limit >= 0), + token_unlimited BOOLEAN NOT NULL DEFAULT false, + cost_limit_usd NUMERIC(12,4) CHECK (cost_limit_usd >= 0), + cost_unlimited BOOLEAN NOT NULL DEFAULT false, + enabled BOOLEAN NOT NULL DEFAULT true, + note TEXT, + created_by TEXT, + updated_by TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CONSTRAINT quota_policies_subject_chk + CHECK ((scope = 'instance') = (subject_id IS NULL)), + CONSTRAINT quota_policies_token_chk + CHECK (NOT (token_unlimited AND token_limit IS NOT NULL)), + CONSTRAINT quota_policies_cost_chk + CHECK (NOT (cost_unlimited AND cost_limit_usd IS NOT NULL)) + ); + """ + ) + # One row per (layer subject, bucket); the instance row's NULL subject folds to ''. + op.execute( + "CREATE UNIQUE INDEX IF NOT EXISTS quota_policies_subject_uidx " + "ON quota_policies (scope, COALESCE(subject_id, ''), bucket);" + ) + op.execute("DROP TRIGGER IF EXISTS quota_policies_set_updated_at ON quota_policies;") + op.execute( + """ + CREATE TRIGGER quota_policies_set_updated_at + BEFORE UPDATE ON quota_policies + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); + """ + ) + op.execute( + """ + CREATE OR REPLACE FUNCTION cleanup_team_quota_policies() RETURNS trigger AS $$ + BEGIN + DELETE FROM quota_policies WHERE scope = 'team' AND subject_id = OLD.id::text; + RETURN OLD; + END; + $$ LANGUAGE plpgsql; + """ + ) + op.execute("DROP TRIGGER IF EXISTS teams_cleanup_quota_policies ON teams;") + op.execute( + """ + CREATE TRIGGER teams_cleanup_quota_policies + AFTER DELETE ON teams + FOR EACH ROW EXECUTE FUNCTION cleanup_team_quota_policies(); + """ + ) + + +def downgrade() -> None: + op.execute("DROP TRIGGER IF EXISTS teams_cleanup_quota_policies ON teams;") + op.execute("DROP FUNCTION IF EXISTS cleanup_team_quota_policies();") + op.execute("DROP TABLE IF EXISTS quota_policies;") + op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cost;") diff --git a/docsgpt/llm/llm_creator.py b/docsgpt/llm/llm_creator.py index 3b3c2a03..995600cb 100644 --- a/docsgpt/llm/llm_creator.py +++ b/docsgpt/llm/llm_creator.py @@ -52,6 +52,7 @@ class LLMCreator: base_url = None upstream_model_id = model_id capabilities = None + model = None if model_id: user_id = model_user_id if user_id is None: @@ -127,4 +128,6 @@ class LLMCreator: # llm.model_id is the upstream name (BYOM resolves it above); stamp # the canonical id (UUID for BYOM) separately for token_usage. llm._canonical_model_id = model_id + # Calls to a user's own model are recorded at $0 (see ``docsgpt/usage.py``). + llm._is_byom = model is not None and getattr(model, "source", "builtin") == "user" return llm diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index a0df1a06..1582e32d 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -28,6 +28,7 @@ from sqlalchemy import ( Index, Integer, MetaData, + Numeric, PrimaryKeyConstraint, UniqueConstraint, Table, @@ -241,6 +242,9 @@ token_usage_table = Table( # cache activity, so hit-rate queries stay honest across providers. Column("cached_tokens", Integer), Column("cache_write_tokens", Integer), + # Added in ``0033_quotas``. USD cost of the call at write time; 0 for + # unpriced and bring-your-own models. + Column("cost", Numeric(12, 8), nullable=False, server_default="0"), ) user_logs_table = Table( @@ -1121,3 +1125,42 @@ Index( personal_access_tokens_table.c.user_id, personal_access_tokens_table.c.created_at.desc(), ) + +# --- Usage quotas (migration 0033) ------------------------------------------ +# Admin-set limits at three layers: instance (``subject_id`` NULL), team +# per-member allowance (``teams.id``) and user override (auth ``sub``). Per +# budget a row sets a limit, marks it unlimited, or defers to the next layer. + +quota_policies_table = Table( + "quota_policies", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()), + Column("scope", Text, nullable=False), + Column("subject_id", Text), + Column("bucket", Text, nullable=False, server_default="all"), + Column("token_limit", BigInteger), + Column("token_unlimited", Boolean, nullable=False, server_default="false"), + Column("cost_limit_usd", Numeric(12, 4)), + Column("cost_unlimited", Boolean, nullable=False, server_default="false"), + Column("enabled", Boolean, nullable=False, server_default="true"), + Column("note", Text), + Column("created_by", Text), + Column("updated_by", Text), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + CheckConstraint("scope IN ('instance', 'team', 'user')", name="quota_policies_scope_check"), + CheckConstraint("bucket IN ('all', 'direct', 'agent')", name="quota_policies_bucket_check"), + CheckConstraint("token_limit >= 0", name="quota_policies_token_limit_check"), + CheckConstraint("cost_limit_usd >= 0", name="quota_policies_cost_limit_check"), + CheckConstraint("(scope = 'instance') = (subject_id IS NULL)", name="quota_policies_subject_chk"), + CheckConstraint("NOT (token_unlimited AND token_limit IS NOT NULL)", name="quota_policies_token_chk"), + CheckConstraint("NOT (cost_unlimited AND cost_limit_usd IS NOT NULL)", name="quota_policies_cost_chk"), +) + +Index( + "quota_policies_subject_uidx", + quota_policies_table.c.scope, + func.coalesce(quota_policies_table.c.subject_id, ""), + quota_policies_table.c.bucket, + unique=True, +) diff --git a/docsgpt/storage/db/repositories/token_usage.py b/docsgpt/storage/db/repositories/token_usage.py index 7292dbd1..83c6e041 100644 --- a/docsgpt/storage/db/repositories/token_usage.py +++ b/docsgpt/storage/db/repositories/token_usage.py @@ -37,6 +37,7 @@ class TokenUsageRepository: timestamp: Optional[datetime] = None, cached_tokens: Optional[int] = None, cache_write_tokens: Optional[int] = None, + cost: float = 0.0, ) -> None: # Attribution guard: the ``token_usage_attribution_chk`` CHECK # constraint requires at least one of ``user_id`` / ``api_key`` @@ -62,14 +63,14 @@ class TokenUsageRepository: INSERT INTO token_usage ( user_id, api_key, agent_id, prompt_tokens, generated_tokens, - cached_tokens, cache_write_tokens, + cached_tokens, cache_write_tokens, cost, source, request_id, model_id, timestamp ) VALUES ( :user_id, :api_key, CAST(:agent_id AS uuid), :prompt_tokens, :generated_tokens, - :cached_tokens, :cache_write_tokens, + :cached_tokens, :cache_write_tokens, :cost, :source, :request_id, :model_id, COALESCE(:timestamp, now()) ) """ @@ -82,6 +83,7 @@ class TokenUsageRepository: "generated_tokens": generated_tokens, "cached_tokens": cached_tokens, "cache_write_tokens": cache_write_tokens, + "cost": cost, "source": source, "request_id": request_id, "model_id": model_id, diff --git a/docsgpt/usage.py b/docsgpt/usage.py index 779163a3..82fc3eb8 100644 --- a/docsgpt/usage.py +++ b/docsgpt/usage.py @@ -2,6 +2,7 @@ import logging import time from typing import Any, Dict +from docsgpt.pricing import compute_cost_usd from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository from docsgpt.storage.db.session import db_session from docsgpt.utils import num_tokens_from_object_or_list, num_tokens_from_string @@ -119,6 +120,12 @@ def _persist_call_usage(llm, call_usage): }, ) return + model_id = getattr(llm, "_canonical_model_id", None) + # Bring-your-own models run on the user's own provider key: recorded, never priced. + if getattr(llm, "_is_byom", False): + cost = 0.0 + else: + cost = _call_cost_usd(model_id, call_usage) try: with db_session() as conn: # ``timestamp`` is omitted so Postgres ``server_default @@ -136,16 +143,32 @@ def _persist_call_usage(llm, call_usage): # "0% cache hits". cached_tokens=call_usage.get("cached_tokens"), cache_write_tokens=call_usage.get("cache_write_tokens"), + cost=cost, source=( getattr(llm, "_token_usage_source", None) or "agent_stream" ), request_id=getattr(llm, "_request_id", None), - model_id=getattr(llm, "_canonical_model_id", None), + model_id=model_id, ) except Exception: logger.exception("token_usage persist failed") +def _call_cost_usd(model_id, call_usage) -> float: + """Price one call; a pricing failure records $0 rather than dropping the row.""" + try: + return compute_cost_usd( + model_id, + call_usage["prompt_tokens"], + call_usage["generated_tokens"], + cached_tokens=call_usage.get("cached_tokens"), + cache_write_tokens=call_usage.get("cache_write_tokens"), + ) + except Exception: + logger.exception("token_usage cost computation failed") + return 0.0 + + def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, int]: """Replace estimates with upstream counts when a provider reported them. diff --git a/tests/llm/test_fallback.py b/tests/llm/test_fallback.py index ceba9b65..03974bb5 100644 --- a/tests/llm/test_fallback.py +++ b/tests/llm/test_fallback.py @@ -1015,6 +1015,38 @@ class TestLLMCreatorPassesModelUserId: assert captured["model_user_id"] == "owner-alice" + @pytest.mark.parametrize( + "model_id, source, expected", + [(None, None, False), ("catalog-model", "builtin", False), ("byom-uuid", "user", True)], + ) + def test_byom_flag_follows_the_model_source(self, monkeypatch, model_id, source, expected): + from types import SimpleNamespace + + from docsgpt.llm.llm_creator import LLMCreator + from docsgpt.llm.providers import PROVIDERS_BY_NAME + + class _LLM: + def __init__(self, *args, **kwargs): + pass + + monkeypatch.setattr(PROVIDERS_BY_NAME["openai"], "llm_class", _LLM) + model = SimpleNamespace( + source=source, api_key="own-key", base_url=None, upstream_model_id=None, capabilities=None + ) + registry = SimpleNamespace(get_model=lambda _id, user_id=None: model) + monkeypatch.setattr( + "docsgpt.core.model_registry.ModelRegistry.get_instance", lambda: registry + ) + + llm = LLMCreator.create_llm( + type="openai", api_key="k", user_api_key=None, + decoded_token={"sub": "u1"}, model_id=model_id, + ) + + assert llm._is_byom is expected + assert llm._canonical_model_id == model_id + + # Tests — responding-provider tracking (cross-provider fallback handler fix) diff --git a/tests/storage/db/repositories/test_token_usage.py b/tests/storage/db/repositories/test_token_usage.py index 46f59d11..02b2eb98 100644 --- a/tests/storage/db/repositories/test_token_usage.py +++ b/tests/storage/db/repositories/test_token_usage.py @@ -35,6 +35,15 @@ class TestInsert: ) assert total == 30 + def test_cost_defaults_to_zero_and_round_trips(self, pg_conn): + repo = _repo(pg_conn) + repo.insert(user_id="u-cost", prompt_tokens=1, generated_tokens=1) + repo.insert(user_id="u-cost", prompt_tokens=1, generated_tokens=1, cost=0.00012345) + costs = pg_conn.execute( + text("SELECT cost FROM token_usage WHERE user_id = 'u-cost' ORDER BY id") + ).scalars().all() + assert [float(c) for c in costs] == [0.0, 0.00012345] + class TestReassignApiKey: def test_rewrites_and_preserves_rate_limit_window(self, pg_conn): diff --git a/tests/storage/db/test_migration_0033.py b/tests/storage/db/test_migration_0033.py new file mode 100644 index 00000000..09cb1d4f --- /dev/null +++ b/tests/storage/db/test_migration_0033.py @@ -0,0 +1,145 @@ +"""Migration round-trip test for 0033_quotas.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError + + +pytestmark = pytest.mark.integration + + +def _alembic_ini() -> Path: + return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini" + + +def _run_alembic(url: str, *args: str) -> None: + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _alembic_heads(url: str) -> list[str]: + out = subprocess.check_output( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + text=True, + ) + return [line for line in out.splitlines() if line.strip()] + + +def _alembic_version(conn) -> str: + return conn.execute(text("SELECT version_num FROM alembic_version")).scalar() + + +def _column_exists(conn, table: str, column: str) -> bool: + row = conn.execute( + text( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = :t AND column_name = :c AND table_schema = 'public'" + ), + {"t": table, "c": column}, + ).fetchone() + return row is not None + + +_0033 = "0033_quotas" +_0032 = "0032_personal_access_tokens" + + +def _table_exists(conn, table: str) -> bool: + return conn.execute(text("SELECT to_regclass(:t)"), {"t": f"public.{table}"}).scalar() is not None + + +def _insert_policy(conn, **values) -> None: + cols = ", ".join(values) + params = ", ".join(f":{k}" for k in values) + conn.execute(text(f"INSERT INTO quota_policies ({cols}) VALUES ({params})"), values) + + +class TestMigration0033RoundTrip: + def test_single_head(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + assert len(_alembic_heads(url)) == 1 + + def test_head_has_quota_schema(self, pg_engine): + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0033 + assert _table_exists(conn, "quota_policies") + assert _column_exists(conn, "token_usage", "cost") + + def test_downgrade_drops_then_upgrade_restores(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0032) + with pg_engine.connect() as conn: + assert _alembic_version(conn) == _0032 + assert not _table_exists(conn, "quota_policies") + assert not _column_exists(conn, "token_usage", "cost") + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + assert _table_exists(conn, "quota_policies") + assert _column_exists(conn, "token_usage", "cost") + + def test_upgrade_tolerates_an_existing_cost_column(self, pg_engine): + """A database that already carries ``token_usage.cost`` upgrades cleanly.""" + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0032) + with pg_engine.begin() as conn: + conn.execute(text("ALTER TABLE token_usage ADD COLUMN cost NUMERIC(12,8) NOT NULL DEFAULT 0")) + conn.execute( + text( + "INSERT INTO token_usage (user_id, prompt_tokens, generated_tokens, cost) " + "VALUES ('u-mig33', 10, 1, 0.5)" + ) + ) + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + cost = conn.execute(text("SELECT cost FROM token_usage WHERE user_id = 'u-mig33'")).scalar() + assert float(cost) == 0.5 + + +class TestQuotaPolicyConstraints: + def test_one_row_per_subject_and_bucket(self, pg_conn): + _insert_policy(pg_conn, scope="instance", token_limit=10) + with pytest.raises(IntegrityError): + with pg_conn.begin_nested(): + _insert_policy(pg_conn, scope="instance", token_limit=20) + _insert_policy(pg_conn, scope="instance", bucket="agent", token_limit=20) + + @pytest.mark.parametrize( + "values", + [ + {"scope": "instance", "subject_id": "u1"}, + {"scope": "user"}, + {"scope": "user", "subject_id": "u1", "token_limit": 5, "token_unlimited": True}, + {"scope": "user", "subject_id": "u1", "cost_limit_usd": 5, "cost_unlimited": True}, + {"scope": "user", "subject_id": "u1", "token_limit": -1}, + {"scope": "user", "subject_id": "u1", "bucket": "nope"}, + {"scope": "org", "subject_id": "u1"}, + ], + ) + def test_invalid_rows_rejected(self, pg_conn, values): + with pytest.raises(IntegrityError): + with pg_conn.begin_nested(): + _insert_policy(pg_conn, **values) + + def test_deleting_a_team_removes_its_policies(self, pg_conn): + team_id = pg_conn.execute( + text("INSERT INTO teams (name, slug, owner_id) VALUES ('Q', 'q-mig33', 'owner') RETURNING id") + ).scalar() + _insert_policy(pg_conn, scope="team", subject_id=str(team_id), token_limit=10) + _insert_policy(pg_conn, scope="user", subject_id=str(team_id), token_limit=10) + pg_conn.execute(text("DELETE FROM teams WHERE id = :id"), {"id": team_id}) + scopes = pg_conn.execute( + text("SELECT scope FROM quota_policies WHERE subject_id = :id"), {"id": str(team_id)} + ).scalars().all() + assert scopes == ["user"] diff --git a/tests/test_usage.py b/tests/test_usage.py index b75639d6..67fbd4b0 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -712,3 +712,56 @@ def test_decorator_omits_cache_bins_when_provider_reports_none(monkeypatch): assert row["cache_write_tokens"] is None assert llm.emitted[0]["cached_tokens"] is None assert llm.emitted[0]["cache_write_tokens"] is None + + +def _persist_with_cost(monkeypatch, llm, cost_fn): + from docsgpt.usage import _persist_call_usage + + _install_fake_token_repo(monkeypatch) + monkeypatch.setattr("docsgpt.usage.compute_cost_usd", cost_fn) + _persist_call_usage( + llm, {"prompt_tokens": 1000, "generated_tokens": 10, "cached_tokens": 400} + ) + return _FakeTokenUsageRepo.last_instance.inserted[0] + + +class _CostLLM: + decoded_token = {"sub": "user_123"} + user_api_key = None + agent_id = None + _canonical_model_id = "priced-model" + + +@pytest.mark.unit +def test_persist_prices_the_call_by_canonical_model(monkeypatch): + seen = {} + + def cost_fn(model, prompt, generated, cached_tokens=None, cache_write_tokens=None): + seen.update(model=model, prompt=prompt, generated=generated, cached=cached_tokens) + return 0.0123 + + row = _persist_with_cost(monkeypatch, _CostLLM(), cost_fn) + + assert row["cost"] == 0.0123 + assert seen == {"model": "priced-model", "prompt": 1000, "generated": 10, "cached": 400} + + +@pytest.mark.unit +def test_persist_records_byom_calls_at_zero_cost(monkeypatch): + llm = _CostLLM() + llm._is_byom = True + + row = _persist_with_cost(monkeypatch, llm, lambda *a, **k: 9.9) + + assert row["cost"] == 0.0 + assert row["prompt_tokens"] == 1000 + + +@pytest.mark.unit +def test_persist_keeps_the_row_when_pricing_fails(monkeypatch): + def boom(*args, **kwargs): + raise RuntimeError("registry unavailable") + + row = _persist_with_cost(monkeypatch, _CostLLM(), boom) + + assert row["cost"] == 0.0 From 6c421392245c189ccbca3e40ed7f76d78346e2bf Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 11:42:53 +0100 Subject: [PATCH 03/14] fix(usage): stop double counting scheduled runs; attribute workflow node usage sum_tokens_in_range now skips the scheduler's per-run rollup rows, whose tokens are already on the run's per-call rows, so the per-agent 24h limit and the admin total no longer count scheduled spend twice. Workflow node LLMs carry the workflow agent's id, so their usage rows are attributed to the agent instead of landing with a user id only. usage_totals returns a user's tokens and cost since a window start, split by interactive and agent-key traffic. --- docsgpt/agents/workflows/workflow_engine.py | 2 + .../storage/db/repositories/token_usage.py | 37 +++++++++++++++-- tests/agents/test_workflow_agent_types.py | 21 ++++++++++ .../db/repositories/test_token_usage.py | 40 +++++++++++++++++++ 4 files changed, 97 insertions(+), 3 deletions(-) diff --git a/docsgpt/agents/workflows/workflow_engine.py b/docsgpt/agents/workflows/workflow_engine.py index 54fa6e67..15876c82 100644 --- a/docsgpt/agents/workflows/workflow_engine.py +++ b/docsgpt/agents/workflows/workflow_engine.py @@ -393,6 +393,8 @@ class WorkflowEngine: "prompt": node_prompt, "chat_history": self.agent.chat_history, "decoded_token": self.agent.decoded_token, + # Attributes the node's token usage to the workflow agent. + "agent_id": getattr(self.agent, "agent_id", None), "json_schema": node_json_schema, "retrieved_docs": node_docs, # A template that interpolates the documents itself already carries diff --git a/docsgpt/storage/db/repositories/token_usage.py b/docsgpt/storage/db/repositories/token_usage.py index 83c6e041..17e546fe 100644 --- a/docsgpt/storage/db/repositories/token_usage.py +++ b/docsgpt/storage/db/repositories/token_usage.py @@ -118,9 +118,13 @@ class TokenUsageRepository: user_id: Optional[str] = None, api_key: Optional[str] = None, ) -> int: - """Total (prompt + generated) tokens in the given time range.""" - clauses = ["timestamp >= :start", "timestamp <= :end"] - params: dict = {"start": start, "end": end} + """Total (prompt + generated) tokens in the given time range. + + Run-level rollup rows (``ROLLUP_SOURCES``) are excluded: their tokens + are already counted on the run's per-call rows. + """ + clauses = ["timestamp >= :start", "timestamp <= :end", "source <> ALL(:rollup_sources)"] + params: dict = {"start": start, "end": end, "rollup_sources": list(self.ROLLUP_SOURCES)} if user_id is not None: clauses.append("user_id = :user_id") params["user_id"] = user_id @@ -134,6 +138,33 @@ class TokenUsageRepository: ) return result.scalar() + def usage_totals(self, *, user_id: str, start: datetime, bucket: str = "all") -> tuple[int, float]: + """Return ``(tokens, cost_usd)`` a user has consumed since ``start``. + + Args: + user_id: The billable user (auth ``sub``). + start: Inclusive window start. + bucket: ``all``, ``direct`` (rows without an agent key) or + ``agent`` (rows with one). + + Rollup rows are excluded; side-channel calls count, they are real spend. + """ + clauses = ["user_id = :user_id", "timestamp >= :start", "source <> ALL(:rollup_sources)"] + if bucket == "direct": + clauses.append("api_key IS NULL") + elif bucket == "agent": + clauses.append("api_key IS NOT NULL") + elif bucket != "all": + raise ValueError(f"unknown usage bucket: {bucket!r}") + row = self._conn.execute( + text( + "SELECT COALESCE(SUM(prompt_tokens + generated_tokens), 0), COALESCE(SUM(cost), 0) " + f"FROM token_usage WHERE {' AND '.join(clauses)}" + ), + {"user_id": user_id, "start": start, "rollup_sources": list(self.ROLLUP_SOURCES)}, + ).one() + return int(row[0]), float(row[1]) + # Token usage written outside a user-initiated request (conversation # title generation, history compression, RAG question condensing, # provider fallback). Mirrors the exclusion list in ``count_in_range``. diff --git a/tests/agents/test_workflow_agent_types.py b/tests/agents/test_workflow_agent_types.py index 41bf7cd4..18bf22d3 100644 --- a/tests/agents/test_workflow_agent_types.py +++ b/tests/agents/test_workflow_agent_types.py @@ -214,6 +214,27 @@ class TestWorkflowEngineAgenticNode: assert engine.state["node_agent_agentic_output"] == "agentic answer" assert engine.state["result"] == "agentic answer" + def test_node_usage_is_attributed_to_the_workflow_agent(self, monkeypatch): + engine = create_engine() + engine.agent.agent_id = "11111111-1111-1111-1111-111111111111" + node = create_agent_node(node_id="agent_attr", agent_type="classic") + + captured: Dict[str, Any] = {} + + def capture_create(**kwargs): + captured.update(kwargs) + return StubNodeAgent([{"answer": "ok"}]) + + monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(capture_create)) + monkeypatch.setattr( + "docsgpt.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + assert captured["agent_id"] == "11111111-1111-1111-1111-111111111111" + def test_agentic_node_passes_retriever_config(self, monkeypatch): engine = create_engine() # The node-source authorization gate is exercised separately; diff --git a/tests/storage/db/repositories/test_token_usage.py b/tests/storage/db/repositories/test_token_usage.py index 02b2eb98..9698c610 100644 --- a/tests/storage/db/repositories/test_token_usage.py +++ b/tests/storage/db/repositories/test_token_usage.py @@ -45,6 +45,46 @@ class TestInsert: assert [float(c) for c in costs] == [0.0, 0.00012345] +class TestRollupExclusion: + def test_sum_tokens_ignores_schedule_rollups(self, pg_conn): + repo = _repo(pg_conn) + repo.insert(user_id="u-roll", prompt_tokens=10, generated_tokens=5, source="agent_stream") + repo.insert(user_id="u-roll", prompt_tokens=10, generated_tokens=5, source="schedule") + total = repo.sum_tokens_in_range( + start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u-roll" + ) + assert total == 15 + + +class TestUsageTotals: + def _seed(self, repo): + repo.insert(user_id="u-tot", prompt_tokens=100, generated_tokens=10, cost=0.5) + repo.insert(user_id="u-tot", api_key="k", prompt_tokens=20, generated_tokens=2, cost=0.25) + repo.insert(user_id="u-tot", prompt_tokens=7, generated_tokens=0, cost=0.125, source="title") + repo.insert(user_id="u-tot", prompt_tokens=999, generated_tokens=0, source="schedule") + repo.insert(user_id="u-other", prompt_tokens=999, generated_tokens=0, cost=9) + repo.insert( + user_id="u-tot", prompt_tokens=999, generated_tokens=0, cost=9, + timestamp=_now() - timedelta(days=40), + ) + + @pytest.mark.parametrize( + "bucket, expected", + [("all", (139, 0.875)), ("direct", (117, 0.625)), ("agent", (22, 0.25))], + ) + def test_totals_per_bucket(self, pg_conn, bucket, expected): + repo = _repo(pg_conn) + self._seed(repo) + assert repo.usage_totals(user_id="u-tot", start=_now() - timedelta(days=1), bucket=bucket) == expected + + def test_no_usage_is_zero(self, pg_conn): + assert _repo(pg_conn).usage_totals(user_id="nobody", start=_now() - timedelta(days=1)) == (0, 0.0) + + def test_unknown_bucket_rejected(self, pg_conn): + with pytest.raises(ValueError): + _repo(pg_conn).usage_totals(user_id="u", start=_now(), bucket="nope") + + class TestReassignApiKey: def test_rewrites_and_preserves_rate_limit_window(self, pg_conn): # Rotating an agent key must carry the running 24h usage window over From 6826313b60ea7cd0026bafe0b2c8540bbff354b9 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 11:46:59 +0100 Subject: [PATCH 04/14] feat(quotas): limit resolution and the quota service docsgpt/quotas resolves a user's effective token and cost limits from the policy rows that apply to them: their own override, then the most generous allowance among their teams, then the instance default, then any defaults a registered provider supplies. Each budget resolves on its own, and a team counts once however many memberships the user holds in it. QuotaService compares those limits with the user's token_usage totals over the current QUOTA_PERIOD window (calendar-aligned, UTC, computed at read time). It fails open, and skips the usage query for unlimited users. --- docsgpt/quotas/__init__.py | 24 ++ docsgpt/quotas/providers.py | 51 ++++ docsgpt/quotas/resolver.py | 99 ++++++++ docsgpt/quotas/service.py | 174 +++++++++++++ docsgpt/quotas/windows.py | 37 +++ .../storage/db/repositories/quota_policies.py | 184 ++++++++++++++ tests/quotas/__init__.py | 0 tests/quotas/test_providers.py | 47 ++++ tests/quotas/test_resolver.py | 124 +++++++++ tests/quotas/test_service.py | 237 ++++++++++++++++++ tests/quotas/test_windows.py | 55 ++++ .../db/repositories/test_quota_policies.py | 143 +++++++++++ 12 files changed, 1175 insertions(+) create mode 100644 docsgpt/quotas/__init__.py create mode 100644 docsgpt/quotas/providers.py create mode 100644 docsgpt/quotas/resolver.py create mode 100644 docsgpt/quotas/service.py create mode 100644 docsgpt/quotas/windows.py create mode 100644 docsgpt/storage/db/repositories/quota_policies.py create mode 100644 tests/quotas/__init__.py create mode 100644 tests/quotas/test_providers.py create mode 100644 tests/quotas/test_resolver.py create mode 100644 tests/quotas/test_service.py create mode 100644 tests/quotas/test_windows.py create mode 100644 tests/storage/db/repositories/test_quota_policies.py diff --git a/docsgpt/quotas/__init__.py b/docsgpt/quotas/__init__.py new file mode 100644 index 00000000..ae1b5f4f --- /dev/null +++ b/docsgpt/quotas/__init__.py @@ -0,0 +1,24 @@ +"""Admin-set usage quotas. + +Limits live in ``quota_policies`` at three layers (instance default, team +per-member allowance, user override), each with a token budget and a USD +budget. ``QuotaService`` resolves a user's effective limits and compares them +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.windows import window_bounds + +__all__ = [ + "BucketStatus", + "QuotaDefaultsProvider", + "QuotaExceeded", + "QuotaService", + "ResolvedLimit", + "ResolvedLimits", + "register_defaults_provider", + "resolve_limits", + "window_bounds", +] diff --git a/docsgpt/quotas/providers.py b/docsgpt/quotas/providers.py new file mode 100644 index 00000000..8b38ef97 --- /dev/null +++ b/docsgpt/quotas/providers.py @@ -0,0 +1,51 @@ +"""Extension point for limits that do not come from ``quota_policies``. + +A deployment can register a provider that supplies per-user default policies +(for example from a subscription plan) and adjusts the quota error payload. +Defaults sit below every stored layer: a stored instance, team or user row +with an opinion always wins. +""" + +from __future__ import annotations + +from typing import Mapping + + +class QuotaDefaultsProvider: + """Base provider: no defaults, error payload unchanged.""" + + def default_policies(self, user_id: str) -> list[dict]: + """Return default policy rows for ``user_id``. + + Each row uses the ``quota_policies`` field names (``bucket``, + ``token_limit``, ``token_unlimited``, ``cost_limit_usd``, + ``cost_unlimited``); ``scope`` is set by the caller. + """ + return [] + + def error_payload(self, payload: dict, user_id: str) -> dict: + """Return the payload sent to a client whose quota is exhausted.""" + return payload + + +_provider: QuotaDefaultsProvider = QuotaDefaultsProvider() + + +def register_defaults_provider(provider: QuotaDefaultsProvider) -> None: + """Replace the process-wide defaults provider.""" + global _provider + _provider = provider + + +def get_defaults_provider() -> QuotaDefaultsProvider: + """Return the registered defaults provider.""" + return _provider + + +def default_rows(user_id: str) -> list[dict]: + """Return the provider's defaults for ``user_id`` as ``default``-layer rows.""" + rows: list[dict] = [] + for row in get_defaults_provider().default_policies(user_id) or []: + if isinstance(row, Mapping): + rows.append({"bucket": "all", "enabled": True, **row, "scope": "default", "subject_id": None}) + return rows diff --git a/docsgpt/quotas/resolver.py b/docsgpt/quotas/resolver.py new file mode 100644 index 00000000..0f3ebb2e --- /dev/null +++ b/docsgpt/quotas/resolver.py @@ -0,0 +1,99 @@ +"""Resolve the policy rows that apply to a user into effective limits. + +Each budget (tokens, cost) resolves on its own: the user's row wins, then the +most generous of the user's team rows, then the instance row, then the +registered defaults. A row with neither a limit nor the unlimited flag for a +budget has no opinion on it and is skipped. + +Teams resolve to the most generous allowance because team membership is not +controlled by the instance admin: under "most restrictive", any team admin +could throttle a user by adding them to a low-allowance team. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable, Mapping, Optional + +LAYERS = ("user", "team", "instance", "default") + +_FIELDS = {"tokens": ("token_limit", "token_unlimited"), "cost": ("cost_limit_usd", "cost_unlimited")} + + +@dataclass(frozen=True) +class ResolvedLimit: + """One budget's effective limit and the layer it came from. + + ``limit`` is ``None`` when the budget is unlimited. ``source`` is ``None`` + when no layer had an opinion, otherwise one of ``LAYERS``; ``source_id`` is + the team id for a team-sourced limit. + """ + + limit: Optional[float] = None + source: Optional[str] = None + source_id: Optional[str] = None + + @property + def unlimited(self) -> bool: + return self.limit is None + + +@dataclass(frozen=True) +class ResolvedLimits: + """A user's effective token and cost limits for one bucket.""" + + tokens: ResolvedLimit + cost: ResolvedLimit + + @property + def unlimited(self) -> bool: + return self.tokens.unlimited and self.cost.unlimited + + +def _opinion(row: Mapping, budget: str) -> Optional[tuple[bool, Optional[float]]]: + """Return ``(unlimited, limit)`` for a row's budget, or ``None`` if it defers.""" + limit_field, unlimited_field = _FIELDS[budget] + if row.get(unlimited_field): + return True, None + value = row.get(limit_field) + if value is None: + return None + return False, float(value) + + +def _resolve_budget(rows: Iterable[Mapping], budget: str) -> ResolvedLimit: + by_layer: dict[str, list[tuple[Mapping, tuple[bool, Optional[float]]]]] = {} + for row in rows: + opinion = _opinion(row, budget) + if opinion is not None: + by_layer.setdefault(row["scope"], []).append((row, opinion)) + for layer in LAYERS: + candidates = by_layer.get(layer) + if not candidates: + continue + # Most generous first: unlimited, then the larger limit. Only the team + # layer can hold more than one candidate. Ties break on subject id so + # the reported source is stable. + row, (unlimited, limit) = min( + candidates, + key=lambda c: (not c[1][0], -(c[1][1] or 0.0), str(c[0].get("subject_id") or "")), + ) + source_id = str(row["subject_id"]) if layer == "team" else None + return ResolvedLimit(limit=None if unlimited else limit, source=layer, source_id=source_id) + return ResolvedLimit() + + +def resolve_limits(rows: Iterable[Mapping], bucket: str = "all") -> ResolvedLimits: + """Return the effective limits for ``bucket`` from a user's applicable rows. + + Args: + rows: Policy rows that apply to the user (their own, their teams', the + instance's and any provider defaults). Disabled rows and rows for + other buckets are ignored. + bucket: The policy bucket to resolve. + """ + applicable = [r for r in rows if r.get("bucket", "all") == bucket and r.get("enabled", True)] + return ResolvedLimits( + tokens=_resolve_budget(applicable, "tokens"), + cost=_resolve_budget(applicable, "cost"), + ) diff --git a/docsgpt/quotas/service.py b/docsgpt/quotas/service.py new file mode 100644 index 00000000..775f3109 --- /dev/null +++ b/docsgpt/quotas/service.py @@ -0,0 +1,174 @@ +"""Compare a user's usage with their effective limits.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from datetime import datetime +from typing import Optional + +from docsgpt.core.settings import settings +from docsgpt.quotas.providers import default_rows, get_defaults_provider +from docsgpt.quotas.resolver import ResolvedLimit, ResolvedLimits, resolve_limits +from docsgpt.quotas.windows import window_bounds +from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository +from docsgpt.storage.db.session import db_readonly + +logger = logging.getLogger(__name__) + +REQUEST_BUCKETS = ("direct", "agent") + +_UNITS = {"tokens": "tokens", "cost": "USD"} + + +@dataclass(frozen=True) +class BucketStatus: + """A user's limits and usage for one policy bucket in the current window.""" + + bucket: str + limits: ResolvedLimits + tokens_used: int + cost_used: float + resets_at: datetime + + def exceeded_budget(self) -> Optional[str]: + """Return ``tokens`` or ``cost`` when that budget is used up, else ``None``.""" + if not self.limits.tokens.unlimited and self.tokens_used >= self.limits.tokens.limit: + return "tokens" + if not self.limits.cost.unlimited and self.cost_used >= self.limits.cost.limit: + return "cost" + return None + + def to_dict(self) -> dict: + """Return the JSON shape shared by the admin and user quota endpoints.""" + + def budget(limit: ResolvedLimit, used: float) -> dict: + return { + "limit": limit.limit, + "used": used, + "source": limit.source, + "source_id": limit.source_id, + } + + return { + "bucket": self.bucket, + "tokens": budget(self.limits.tokens, self.tokens_used), + "cost": budget(self.limits.cost, round(self.cost_used, 6)), + "resets_at": self.resets_at.isoformat(), + } + + +@dataclass(frozen=True) +class QuotaExceeded: + """The exhausted budget that blocks a request.""" + + user_id: str + bucket: str + budget: str + usage: float + limit: float + source: Optional[str] + source_id: Optional[str] + resets_at: datetime + + @property + def retry_after_seconds(self) -> int: + now = datetime.now(self.resets_at.tzinfo) + return max(int((self.resets_at - now).total_seconds()), 1) + + def to_payload(self) -> dict: + """Return the client-facing error body, after the provider's adjustments.""" + unit = _UNITS[self.budget] + if self.budget == "cost": + amounts = f"${self.usage:.2f} of ${self.limit:.2f}" + else: + amounts = f"{int(self.usage):,} of {int(self.limit):,} tokens" + payload = { + "success": False, + "error_code": "quota-exceeded", + "message": f"Usage quota reached ({amounts}). It resets at {self.resets_at.isoformat()}.", + "limit_scope": "user_quota", + "dimension": self.budget, + "unit": unit, + "usage": round(self.usage, 6) if self.budget == "cost" else int(self.usage), + "limit": self.limit if self.budget == "cost" else int(self.limit), + "bucket": self.bucket, + "source": self.source, + "resets_at": self.resets_at.isoformat(), + } + try: + return get_defaults_provider().error_payload(payload, self.user_id) or payload + except Exception: + logger.exception("quota defaults provider failed to build the error payload") + return payload + + +class QuotaService: + """Resolve limits and measure usage for the current ``QUOTA_PERIOD`` window.""" + + @staticmethod + def status( + user_id: str, + buckets: tuple[str, ...] = ("all",), + now: Optional[datetime] = None, + ) -> list[BucketStatus]: + """Return the user's status for each of ``buckets``. + + Usage is only summed for buckets that carry a limit, so a user with no + applicable policy costs one policy lookup and no usage query. + """ + start, resets_at = window_bounds(settings.QUOTA_PERIOD, now) + statuses: list[BucketStatus] = [] + with db_readonly() as conn: + rows = QuotaPoliciesRepository(conn).policies_for_user(user_id) + default_rows(user_id) + usage_repo = TokenUsageRepository(conn) + for bucket in buckets: + limits = resolve_limits(rows, bucket) + tokens_used, cost_used = (0, 0.0) + if not limits.unlimited: + tokens_used, cost_used = usage_repo.usage_totals(user_id=user_id, start=start, bucket=bucket) + statuses.append(BucketStatus(bucket, limits, tokens_used, cost_used, resets_at)) + return statuses + + @classmethod + def check( + cls, user_id: Optional[str], bucket: str = "direct", now: Optional[datetime] = None + ) -> Optional[QuotaExceeded]: + """Return why ``user_id`` may not start a request, or ``None`` if they may. + + Both the ``all`` policies and the request's own bucket must have room. + The check runs before the request, so the call that crosses a limit + completes and the next one is refused. Any failure here allows the + request: a quota outage must not take chat down. + + Args: + user_id: The billable user. Requests with no user are not limited. + bucket: ``direct`` for chat without an agent, ``agent`` for traffic through one. + now: Reference instant, for tests. + """ + if not user_id: + return None + if bucket not in REQUEST_BUCKETS: + raise ValueError(f"unknown request bucket: {bucket!r}") + try: + statuses = cls.status(user_id, ("all", bucket), now) + except Exception: + logger.exception("quota check failed; allowing the request", extra={"user_id": user_id}) + return None + for status in statuses: + budget = status.exceeded_budget() + if budget is None: + continue + limit = status.limits.tokens if budget == "tokens" else status.limits.cost + return QuotaExceeded( + user_id=user_id, + bucket=status.bucket, + budget=budget, + usage=status.tokens_used if budget == "tokens" else status.cost_used, + limit=limit.limit, + source=limit.source, + source_id=limit.source_id, + resets_at=status.resets_at, + ) + return None diff --git a/docsgpt/quotas/windows.py b/docsgpt/quotas/windows.py new file mode 100644 index 00000000..58f5e67c --- /dev/null +++ b/docsgpt/quotas/windows.py @@ -0,0 +1,37 @@ +"""Calendar-aligned UTC quota windows, computed at read time (no reset job).""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from typing import Optional + +PERIODS = ("day", "week", "month") + + +def window_bounds(period: str, now: Optional[datetime] = None) -> tuple[datetime, datetime]: + """Return ``(start, resets_at)`` of the window containing ``now``. + + Args: + period: ``day`` (from 00:00), ``week`` (from Monday) or ``month`` (from the 1st). + now: Reference instant; defaults to the current time. Naive values are read as UTC. + + Raises: + ValueError: If ``period`` is not one of ``PERIODS``. + """ + if now is None: + now = datetime.now(timezone.utc) + elif now.tzinfo is None: + now = now.replace(tzinfo=timezone.utc) + else: + now = now.astimezone(timezone.utc) + midnight = now.replace(hour=0, minute=0, second=0, microsecond=0) + if period == "day": + return midnight, midnight + timedelta(days=1) + if period == "week": + start = midnight - timedelta(days=midnight.weekday()) + return start, start + timedelta(days=7) + if period == "month": + start = midnight.replace(day=1) + next_month = (start + timedelta(days=32)).replace(day=1) + return start, next_month + raise ValueError(f"unknown quota period: {period!r}") diff --git a/docsgpt/storage/db/repositories/quota_policies.py b/docsgpt/storage/db/repositories/quota_policies.py new file mode 100644 index 00000000..3713fc48 --- /dev/null +++ b/docsgpt/storage/db/repositories/quota_policies.py @@ -0,0 +1,184 @@ +"""Repository for the ``quota_policies`` table. + +One row per ``(scope, subject, bucket)``: the instance default (no subject), a +team's per-member allowance (``teams.id``) or a user's override (auth ``sub``). +All methods take a ``Connection`` and do not manage their own transactions. +""" + +from __future__ import annotations + +from typing import Optional + +from sqlalchemy import Connection, text + +from docsgpt.storage.db.base_repository import row_to_dict + +SCOPES = ("instance", "team", "user") +BUCKETS = ("all", "direct", "agent") + +_COLUMNS = ( + "id, scope, subject_id, bucket, token_limit, token_unlimited, cost_limit_usd, " + "cost_unlimited, enabled, note, created_by, updated_by, created_at, updated_at" +) + + +def _validate(scope: str, subject_id: Optional[str], bucket: str) -> None: + if scope not in SCOPES: + raise ValueError(f"unknown quota scope: {scope!r}") + if bucket not in BUCKETS: + raise ValueError(f"unknown quota bucket: {bucket!r}") + if (scope == "instance") != (subject_id is None): + raise ValueError("subject_id is required for team and user policies and must be omitted for instance") + + +class QuotaPoliciesRepository: + """Admin-set usage limits.""" + + def __init__(self, conn: Connection) -> None: + self._conn = conn + + # ------------------------------------------------------------------ + # Reads + # ------------------------------------------------------------------ + def policies_for_user(self, user_id: str) -> list[dict]: + """Return every enabled row that applies to ``user_id``. + + That is the instance rows, the user's own rows, and the rows of each + team the user belongs to (once per team, whatever roles or sources + the membership has). + """ + result = self._conn.execute( + text( + f""" + SELECT {_COLUMNS} FROM quota_policies + WHERE enabled AND ( + scope = 'instance' + OR (scope = 'user' AND subject_id = :user_id) + OR (scope = 'team' AND subject_id IN ( + SELECT DISTINCT team_id::text FROM team_members WHERE user_id = :user_id + )) + ) + """ + ), + {"user_id": user_id}, + ) + return [row_to_dict(row) for row in result.fetchall()] + + def get(self, scope: str, subject_id: Optional[str], bucket: str = "all") -> Optional[dict]: + """Return one policy row, or ``None``.""" + _validate(scope, subject_id, bucket) + row = self._conn.execute( + text( + f"SELECT {_COLUMNS} FROM quota_policies " + "WHERE scope = :scope AND COALESCE(subject_id, '') = :subject AND bucket = :bucket" + ), + {"scope": scope, "subject": subject_id or "", "bucket": bucket}, + ).fetchone() + return row_to_dict(row) if row is not None else None + + def list_for_subject(self, scope: str, subject_id: Optional[str]) -> list[dict]: + """Return a subject's rows across buckets, ``all`` first.""" + _validate(scope, subject_id, "all") + result = self._conn.execute( + text( + f"SELECT {_COLUMNS} FROM quota_policies " + "WHERE scope = :scope AND COALESCE(subject_id, '') = :subject " + "ORDER BY array_position(ARRAY['all', 'direct', 'agent'], bucket)" + ), + {"scope": scope, "subject": subject_id or ""}, + ) + return [row_to_dict(row) for row in result.fetchall()] + + def list_by_scope(self, scope: str) -> list[dict]: + """Return every row of a scope, ordered by subject then bucket.""" + if scope not in SCOPES: + raise ValueError(f"unknown quota scope: {scope!r}") + result = self._conn.execute( + text( + f"SELECT {_COLUMNS} FROM quota_policies WHERE scope = :scope " + "ORDER BY subject_id NULLS FIRST, array_position(ARRAY['all', 'direct', 'agent'], bucket)" + ), + {"scope": scope}, + ) + return [row_to_dict(row) for row in result.fetchall()] + + # ------------------------------------------------------------------ + # Writes + # ------------------------------------------------------------------ + def upsert( + self, + *, + scope: str, + subject_id: Optional[str], + bucket: str = "all", + token_limit: Optional[int] = None, + token_unlimited: bool = False, + cost_limit_usd: Optional[float] = None, + cost_unlimited: bool = False, + enabled: bool = True, + note: Optional[str] = None, + actor: Optional[str] = None, + ) -> dict: + """Create or replace the policy for ``(scope, subject_id, bucket)``. + + Raises: + ValueError: On an unknown scope or bucket, a subject that does not + match the scope, a negative limit, or a budget that is both + limited and unlimited. + """ + _validate(scope, subject_id, bucket) + if token_unlimited and token_limit is not None: + raise ValueError("token budget cannot be both limited and unlimited") + if cost_unlimited and cost_limit_usd is not None: + raise ValueError("cost budget cannot be both limited and unlimited") + if (token_limit is not None and token_limit < 0) or (cost_limit_usd is not None and cost_limit_usd < 0): + raise ValueError("limits must not be negative") + row = self._conn.execute( + text( + f""" + INSERT INTO quota_policies ( + scope, subject_id, bucket, token_limit, token_unlimited, + cost_limit_usd, cost_unlimited, enabled, note, created_by, updated_by + ) + VALUES ( + :scope, :subject_id, :bucket, :token_limit, :token_unlimited, + :cost_limit_usd, :cost_unlimited, :enabled, :note, :actor, :actor + ) + ON CONFLICT (scope, COALESCE(subject_id, ''), bucket) DO UPDATE SET + token_limit = EXCLUDED.token_limit, + token_unlimited = EXCLUDED.token_unlimited, + cost_limit_usd = EXCLUDED.cost_limit_usd, + cost_unlimited = EXCLUDED.cost_unlimited, + enabled = EXCLUDED.enabled, + note = EXCLUDED.note, + updated_by = EXCLUDED.updated_by + RETURNING {_COLUMNS} + """ + ), + { + "scope": scope, + "subject_id": subject_id, + "bucket": bucket, + "token_limit": token_limit, + "token_unlimited": token_unlimited, + "cost_limit_usd": cost_limit_usd, + "cost_unlimited": cost_unlimited, + "enabled": enabled, + "note": note, + "actor": actor, + }, + ).one() + return row_to_dict(row) + + def delete(self, scope: str, subject_id: Optional[str], bucket: Optional[str] = None) -> int: + """Delete a subject's policy for ``bucket``, or all of them; return the count.""" + _validate(scope, subject_id, bucket or "all") + clauses = ["scope = :scope", "COALESCE(subject_id, '') = :subject"] + params = {"scope": scope, "subject": subject_id or ""} + if bucket is not None: + clauses.append("bucket = :bucket") + params["bucket"] = bucket + result = self._conn.execute( + text(f"DELETE FROM quota_policies WHERE {' AND '.join(clauses)}"), params + ) + return result.rowcount diff --git a/tests/quotas/__init__.py b/tests/quotas/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/quotas/test_providers.py b/tests/quotas/test_providers.py new file mode 100644 index 00000000..aa51383c --- /dev/null +++ b/tests/quotas/test_providers.py @@ -0,0 +1,47 @@ +"""Tests for docsgpt/quotas/providers.py.""" + +from __future__ import annotations + +import pytest + +from docsgpt.quotas import providers +from docsgpt.quotas.providers import QuotaDefaultsProvider, default_rows, register_defaults_provider +from docsgpt.quotas.resolver import resolve_limits + + +@pytest.fixture(autouse=True) +def _restore_provider(): + original = providers.get_defaults_provider() + yield + register_defaults_provider(original) + + +class _PlanProvider(QuotaDefaultsProvider): + def default_policies(self, user_id): + return [ + {"bucket": "agent", "cost_limit_usd": 5.0, "scope": "user"}, + {"cost_limit_usd": 10.0}, + "ignored", + ] + + def error_payload(self, payload, user_id): + return {**payload, "error_code": "free-limit-reached"} + + +@pytest.mark.unit +class TestDefaultsProvider: + def test_base_provider_has_no_defaults(self): + assert default_rows("u1") == [] + assert QuotaDefaultsProvider().error_payload({"a": 1}, "u1") == {"a": 1} + + def test_rows_are_forced_into_the_default_layer(self): + register_defaults_provider(_PlanProvider()) + rows = default_rows("u1") + assert [r["scope"] for r in rows] == ["default", "default"] + assert [r["bucket"] for r in rows] == ["agent", "all"] + assert resolve_limits(rows, "agent").cost.source == "default" + + def test_stored_rows_win_over_defaults(self): + register_defaults_provider(_PlanProvider()) + stored = {"scope": "instance", "subject_id": None, "bucket": "all", "enabled": True, "cost_limit_usd": 99} + assert resolve_limits(default_rows("u1") + [stored]).cost.limit == 99.0 diff --git a/tests/quotas/test_resolver.py b/tests/quotas/test_resolver.py new file mode 100644 index 00000000..4401f0b7 --- /dev/null +++ b/tests/quotas/test_resolver.py @@ -0,0 +1,124 @@ +"""Tests for docsgpt/quotas/resolver.py.""" + +from __future__ import annotations + +import pytest + +from docsgpt.quotas.resolver import ResolvedLimit, resolve_limits + + +def _row(scope, subject_id=None, **fields): + return {"scope": scope, "subject_id": subject_id, "bucket": "all", "enabled": True, **fields} + + +@pytest.mark.unit +class TestLayers: + def test_no_rows_is_unlimited(self): + limits = resolve_limits([]) + assert limits.unlimited + assert limits.tokens == ResolvedLimit() + + def test_user_beats_team_beats_instance_beats_default(self): + rows = [ + _row("default", token_limit=1), + _row("instance", token_limit=10), + _row("team", "t1", token_limit=100), + _row("user", "u1", token_limit=5), + ] + assert resolve_limits(rows).tokens == ResolvedLimit(5.0, "user") + assert resolve_limits(rows[:3]).tokens == ResolvedLimit(100.0, "team", "t1") + assert resolve_limits(rows[:2]).tokens == ResolvedLimit(10.0, "instance") + assert resolve_limits(rows[:1]).tokens == ResolvedLimit(1.0, "default") + + def test_user_override_can_be_stricter_than_the_team(self): + rows = [_row("team", "t1", token_unlimited=True), _row("user", "u1", token_limit=0)] + assert resolve_limits(rows).tokens == ResolvedLimit(0.0, "user") + + def test_user_unlimited_lifts_an_instance_limit(self): + rows = [_row("instance", token_limit=10), _row("user", "u1", token_unlimited=True)] + resolved = resolve_limits(rows).tokens + assert resolved.unlimited and resolved.source == "user" + + def test_budgets_resolve_independently(self): + rows = [ + _row("instance", token_limit=10, cost_limit_usd=1), + _row("user", "u1", cost_limit_usd=25), + ] + limits = resolve_limits(rows) + assert limits.tokens == ResolvedLimit(10.0, "instance") + assert limits.cost == ResolvedLimit(25.0, "user") + + def test_a_row_with_no_opinion_defers(self): + rows = [_row("instance", token_limit=10), _row("user", "u1", note="vip")] + assert resolve_limits(rows).tokens == ResolvedLimit(10.0, "instance") + + def test_zero_is_a_limit_not_unlimited(self): + resolved = resolve_limits([_row("instance", cost_limit_usd=0)]).cost + assert resolved.limit == 0.0 and not resolved.unlimited + + +@pytest.mark.unit +class TestMultipleTeams: + def test_most_generous_team_wins(self): + rows = [ + _row("team", "small", token_limit=100), + _row("team", "big", token_limit=900), + _row("team", "mid", token_limit=500), + ] + assert resolve_limits(rows).tokens == ResolvedLimit(900.0, "team", "big") + + def test_an_unlimited_team_beats_any_limit(self): + rows = [_row("team", "big", token_limit=10**12), _row("team", "free", token_unlimited=True)] + assert resolve_limits(rows).tokens == ResolvedLimit(None, "team", "free") + + def test_allowances_are_not_added_together(self): + rows = [_row("team", "a", token_limit=100), _row("team", "b", token_limit=100)] + assert resolve_limits(rows).tokens.limit == 100.0 + + def test_equal_teams_report_a_stable_source(self): + rows = [_row("team", "b", token_limit=100), _row("team", "a", token_limit=100)] + assert resolve_limits(rows).tokens.source_id == "a" + assert resolve_limits(list(reversed(rows))).tokens.source_id == "a" + + def test_each_budget_can_come_from_a_different_team(self): + rows = [ + _row("team", "tok", token_limit=900, cost_limit_usd=1), + _row("team", "usd", token_limit=100, cost_limit_usd=50), + ] + limits = resolve_limits(rows) + assert limits.tokens.source_id == "tok" + assert limits.cost.source_id == "usd" + + def test_a_team_without_an_opinion_does_not_lift_the_limit(self): + rows = [_row("team", "quiet"), _row("team", "capped", token_limit=100)] + assert resolve_limits(rows).tokens == ResolvedLimit(100.0, "team", "capped") + + def test_teams_without_opinions_fall_through_to_instance(self): + rows = [_row("team", "quiet", cost_limit_usd=5), _row("instance", token_limit=10)] + assert resolve_limits(rows).tokens == ResolvedLimit(10.0, "instance") + + def test_a_zero_team_does_not_block_a_member_of_a_funded_team(self): + rows = [_row("team", "blocked", token_limit=0), _row("team", "funded", token_limit=50)] + assert resolve_limits(rows).tokens.limit == 50.0 + + def test_disabled_team_rows_are_ignored(self): + rows = [_row("team", "big", token_limit=900, enabled=False), _row("team", "small", token_limit=100)] + assert resolve_limits(rows).tokens == ResolvedLimit(100.0, "team", "small") + + +@pytest.mark.unit +class TestBuckets: + def test_only_rows_of_the_bucket_apply(self): + rows = [ + _row("instance", token_limit=10), + {**_row("instance", token_limit=3), "bucket": "agent"}, + ] + assert resolve_limits(rows, "all").tokens.limit == 10.0 + assert resolve_limits(rows, "agent").tokens.limit == 3.0 + assert resolve_limits(rows, "direct").unlimited + + def test_decimal_limits_become_floats(self): + from decimal import Decimal + + resolved = resolve_limits([_row("instance", cost_limit_usd=Decimal("12.5000"))]).cost + assert resolved.limit == 12.5 and isinstance(resolved.limit, float) diff --git a/tests/quotas/test_service.py b/tests/quotas/test_service.py new file mode 100644 index 00000000..0a365799 --- /dev/null +++ b/tests/quotas/test_service.py @@ -0,0 +1,237 @@ +"""Tests for QuotaService against a real Postgres instance.""" + +from __future__ import annotations + +from contextlib import contextmanager +from datetime import datetime, timedelta, timezone + +import pytest +from sqlalchemy import text + +from docsgpt.quotas import providers +from docsgpt.quotas.providers import QuotaDefaultsProvider, register_defaults_provider +from docsgpt.quotas.service import QuotaService +from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository + +NOW = datetime(2026, 9, 23, 12, tzinfo=timezone.utc) +THIS_MONTH = NOW - timedelta(days=2) +LAST_MONTH = NOW - timedelta(days=40) + + +@pytest.fixture +def conn(pg_conn, monkeypatch): + @contextmanager + def _readonly(): + yield pg_conn + + monkeypatch.setattr("docsgpt.quotas.service.db_readonly", _readonly) + monkeypatch.setattr("docsgpt.quotas.service.settings.QUOTA_PERIOD", "month") + return pg_conn + + +@pytest.fixture(autouse=True) +def _restore_provider(): + original = providers.get_defaults_provider() + yield + register_defaults_provider(original) + + +def _use(conn, user_id="u1", tokens=0, cost=0.0, api_key=None, when=THIS_MONTH, source="agent_stream"): + TokenUsageRepository(conn).insert( + user_id=user_id, api_key=api_key, prompt_tokens=tokens, cost=cost, timestamp=when, source=source + ) + + +def _policy(conn, scope, subject_id=None, **fields): + return QuotaPoliciesRepository(conn).upsert(scope=scope, subject_id=subject_id, **fields) + + +def _team_with_member(conn, slug, user_id="u1"): + team_id = str( + conn.execute( + text("INSERT INTO teams (name, slug, owner_id) VALUES (:n, :s, 'o') RETURNING id"), + {"n": slug, "s": slug}, + ).scalar() + ) + conn.execute( + text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), :u, 'team_member')"), + {"t": team_id, "u": user_id}, + ) + return team_id + + +class TestCheck: + def test_no_policies_allows(self, conn): + _use(conn, tokens=10**9) + assert QuotaService.check("u1", now=NOW) is None + + def test_no_user_allows(self, conn): + assert QuotaService.check(None, now=NOW) is None + + def test_under_the_limit_allows(self, conn): + _policy(conn, "instance", token_limit=100) + _use(conn, tokens=99) + assert QuotaService.check("u1", now=NOW) is None + + def test_reaching_the_token_limit_blocks(self, conn): + _policy(conn, "instance", token_limit=100) + _use(conn, tokens=100) + exceeded = QuotaService.check("u1", now=NOW) + assert (exceeded.budget, exceeded.usage, exceeded.limit, exceeded.source) == ("tokens", 100, 100.0, "instance") + assert exceeded.resets_at == datetime(2026, 10, 1, tzinfo=timezone.utc) + + def test_cost_limit_blocks(self, conn): + _policy(conn, "user", "u1", cost_limit_usd=1.5) + _use(conn, cost=1.0) + _use(conn, cost=0.5) + exceeded = QuotaService.check("u1", now=NOW) + assert (exceeded.budget, exceeded.usage, exceeded.source) == ("cost", 1.5, "user") + + def test_zero_limit_blocks_without_usage(self, conn): + _policy(conn, "user", "u1", token_limit=0) + assert QuotaService.check("u1", now=NOW).limit == 0 + + def test_last_periods_usage_does_not_count(self, conn): + _policy(conn, "instance", token_limit=100) + _use(conn, tokens=500, when=LAST_MONTH) + assert QuotaService.check("u1", now=NOW) is None + + def test_other_users_usage_does_not_count(self, conn): + _policy(conn, "instance", token_limit=100) + _use(conn, user_id="u2", tokens=500) + assert QuotaService.check("u1", now=NOW) is None + + def test_scheduler_rollups_do_not_count(self, conn): + _policy(conn, "instance", token_limit=100) + _use(conn, tokens=60) + _use(conn, tokens=60, source="schedule") + assert QuotaService.check("u1", now=NOW) is None + + def test_user_override_lifts_the_instance_limit(self, conn): + _policy(conn, "instance", token_limit=100) + _policy(conn, "user", "u1", token_unlimited=True) + _use(conn, tokens=10**6) + assert QuotaService.check("u1", now=NOW) is None + + def test_period_setting_moves_the_window(self, conn, monkeypatch): + monkeypatch.setattr("docsgpt.quotas.service.settings.QUOTA_PERIOD", "day") + _policy(conn, "instance", token_limit=100) + _use(conn, tokens=500, when=NOW - timedelta(days=1)) + assert QuotaService.check("u1", now=NOW) is None + _use(conn, tokens=500, when=NOW - timedelta(hours=1)) + assert QuotaService.check("u1", now=NOW).resets_at == datetime(2026, 9, 24, tzinfo=timezone.utc) + + def test_unknown_bucket_rejected(self, conn): + with pytest.raises(ValueError): + QuotaService.check("u1", bucket="all", now=NOW) + + def test_failure_allows_the_request(self, monkeypatch): + def boom(): + raise RuntimeError("db down") + + monkeypatch.setattr("docsgpt.quotas.service.db_readonly", boom) + assert QuotaService.check("u1", now=NOW) is None + + +class TestTeams: + def test_the_most_generous_team_sets_the_allowance(self, conn): + _policy(conn, "team", _team_with_member(conn, "svc-small"), token_limit=100) + big = _team_with_member(conn, "svc-big") + _policy(conn, "team", big, token_limit=1000) + _use(conn, tokens=500) + assert QuotaService.check("u1", now=NOW) is None + _use(conn, tokens=500) + exceeded = QuotaService.check("u1", now=NOW) + assert (exceeded.limit, exceeded.source, exceeded.source_id) == (1000.0, "team", big) + + def test_usage_is_one_total_not_one_per_team(self, conn): + for slug in ("svc-a", "svc-b", "svc-c"): + _policy(conn, "team", _team_with_member(conn, slug), token_limit=100) + _use(conn, tokens=100) + assert QuotaService.check("u1", now=NOW).limit == 100.0 + + def test_a_team_only_covers_its_members(self, conn): + _policy(conn, "instance", token_limit=100) + _policy(conn, "team", _team_with_member(conn, "svc-vip", user_id="u2"), token_unlimited=True) + _use(conn, tokens=100) + assert QuotaService.check("u1", now=NOW).source == "instance" + + +class TestBuckets: + def test_bucket_policies_only_see_their_traffic(self, conn): + _policy(conn, "instance", bucket="agent", token_limit=100) + _use(conn, tokens=500) + _use(conn, tokens=90, api_key="k") + assert QuotaService.check("u1", "direct", now=NOW) is None + assert QuotaService.check("u1", "agent", now=NOW) is None + _use(conn, tokens=10, api_key="k") + exceeded = QuotaService.check("u1", "agent", now=NOW) + assert (exceeded.bucket, exceeded.usage) == ("agent", 100) + assert QuotaService.check("u1", "direct", now=NOW) is None + + def test_the_all_bucket_applies_to_both_kinds_of_traffic(self, conn): + _policy(conn, "instance", token_limit=100) + _use(conn, tokens=60) + _use(conn, tokens=60, api_key="k") + assert QuotaService.check("u1", "direct", now=NOW).bucket == "all" + assert QuotaService.check("u1", "agent", now=NOW).bucket == "all" + + +class TestProviderDefaults: + class _Plan(QuotaDefaultsProvider): + def default_policies(self, user_id): + return [{"cost_limit_usd": 5.0}] if user_id == "u1" else [] + + def error_payload(self, payload, user_id): + return {**payload, "error_code": "free-limit-reached"} + + def test_defaults_apply_without_stored_rows(self, conn): + register_defaults_provider(self._Plan()) + _use(conn, cost=5.0) + exceeded = QuotaService.check("u1", now=NOW) + assert exceeded.source == "default" + assert exceeded.to_payload()["error_code"] == "free-limit-reached" + + def test_a_broken_provider_payload_falls_back(self, conn): + class _Broken(self._Plan): + def error_payload(self, payload, user_id): + raise RuntimeError("nope") + + register_defaults_provider(_Broken()) + _use(conn, cost=5.0) + assert QuotaService.check("u1", now=NOW).to_payload()["error_code"] == "quota-exceeded" + + +class TestStatusAndPayload: + def test_status_reports_limits_and_usage(self, conn): + _policy(conn, "instance", token_limit=100, cost_limit_usd=2) + _use(conn, tokens=40, cost=0.5) + (status,) = QuotaService.status("u1", now=NOW) + assert status.to_dict() == { + "bucket": "all", + "tokens": {"limit": 100.0, "used": 40, "source": "instance", "source_id": None}, + "cost": {"limit": 2.0, "used": 0.5, "source": "instance", "source_id": None}, + "resets_at": "2026-10-01T00:00:00+00:00", + } + + def test_unlimited_users_skip_the_usage_query(self, conn, monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("usage must not be summed for an unlimited user") + + monkeypatch.setattr(TokenUsageRepository, "usage_totals", fail) + (status,) = QuotaService.status("u1", now=NOW) + assert status.limits.unlimited and status.tokens_used == 0 + + def test_payload_shape(self, conn): + _policy(conn, "user", "u1", cost_limit_usd=1) + _use(conn, cost=1.25) + exceeded = QuotaService.check("u1", now=NOW) + payload = exceeded.to_payload() + assert payload["success"] is False + assert payload["error_code"] == "quota-exceeded" + assert (payload["dimension"], payload["unit"]) == ("cost", "USD") + assert (payload["usage"], payload["limit"]) == (1.25, 1.0) + assert payload["resets_at"] == "2026-10-01T00:00:00+00:00" + assert "$1.25 of $1.00" in payload["message"] + assert exceeded.retry_after_seconds >= 1 diff --git a/tests/quotas/test_windows.py b/tests/quotas/test_windows.py new file mode 100644 index 00000000..2a243ff2 --- /dev/null +++ b/tests/quotas/test_windows.py @@ -0,0 +1,55 @@ +"""Tests for docsgpt/quotas/windows.py.""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone + +import pytest + +from docsgpt.quotas.windows import window_bounds + +UTC = timezone.utc + + +@pytest.mark.unit +class TestWindowBounds: + @pytest.mark.parametrize( + "period, start, end", + [ + ("day", datetime(2026, 9, 23, tzinfo=UTC), datetime(2026, 9, 24, tzinfo=UTC)), + ("week", datetime(2026, 9, 21, tzinfo=UTC), datetime(2026, 9, 28, tzinfo=UTC)), + ("month", datetime(2026, 9, 1, tzinfo=UTC), datetime(2026, 10, 1, tzinfo=UTC)), + ], + ) + def test_midweek(self, period, start, end): + now = datetime(2026, 9, 23, 15, 30, 12, 99, tzinfo=UTC) # a Wednesday + assert window_bounds(period, now) == (start, end) + + def test_week_starts_on_monday_itself(self): + monday = datetime(2026, 9, 21, 0, 0, tzinfo=UTC) + assert window_bounds("week", monday)[0] == monday + + def test_month_rolls_over_the_year(self): + start, end = window_bounds("month", datetime(2026, 12, 31, 23, 59, tzinfo=UTC)) + assert (start, end) == (datetime(2026, 12, 1, tzinfo=UTC), datetime(2027, 1, 1, tzinfo=UTC)) + + def test_leap_february(self): + start, end = window_bounds("month", datetime(2028, 2, 29, 12, tzinfo=UTC)) + assert (end - start).days == 29 + + def test_other_timezones_are_read_in_utc(self): + tokyo = timezone(timedelta(hours=9)) + # 2026-10-01 08:00 in Tokyo is still 2026-09-30 in UTC. + start, _ = window_bounds("month", datetime(2026, 10, 1, 8, tzinfo=tokyo)) + assert start == datetime(2026, 9, 1, tzinfo=UTC) + + def test_naive_datetimes_are_utc(self): + assert window_bounds("day", datetime(2026, 9, 23, 5))[0] == datetime(2026, 9, 23, tzinfo=UTC) + + def test_defaults_to_now(self): + start, end = window_bounds("day") + assert start <= datetime.now(UTC) < end + + def test_unknown_period(self): + with pytest.raises(ValueError): + window_bounds("year") diff --git a/tests/storage/db/repositories/test_quota_policies.py b/tests/storage/db/repositories/test_quota_policies.py new file mode 100644 index 00000000..25396210 --- /dev/null +++ b/tests/storage/db/repositories/test_quota_policies.py @@ -0,0 +1,143 @@ +"""Tests for QuotaPoliciesRepository against a real Postgres instance.""" + +from __future__ import annotations + +import pytest +from sqlalchemy import text + +from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository + + +def _repo(conn) -> QuotaPoliciesRepository: + return QuotaPoliciesRepository(conn) + + +def _team(conn, slug: str) -> str: + return str( + conn.execute( + text("INSERT INTO teams (name, slug, owner_id) VALUES (:n, :s, 'owner') RETURNING id"), + {"n": slug, "s": slug}, + ).scalar() + ) + + +def _member(conn, team_id: str, user_id: str, role: str = "team_member", source: str = "manual") -> None: + conn.execute( + text( + "INSERT INTO team_members (team_id, user_id, role, source) " + "VALUES (CAST(:t AS uuid), :u, :r, :s)" + ), + {"t": team_id, "u": user_id, "r": role, "s": source}, + ) + + +class TestUpsert: + def test_creates_then_replaces(self, pg_conn): + repo = _repo(pg_conn) + created = repo.upsert(scope="user", subject_id="u1", token_limit=100, note="trial", actor="admin1") + assert created["token_limit"] == 100 + assert created["created_by"] == created["updated_by"] == "admin1" + + replaced = repo.upsert(scope="user", subject_id="u1", cost_limit_usd=2.5, actor="admin2") + assert replaced["id"] == created["id"] + assert replaced["token_limit"] is None + assert float(replaced["cost_limit_usd"]) == 2.5 + assert replaced["note"] is None + assert (replaced["created_by"], replaced["updated_by"]) == ("admin1", "admin2") + + def test_instance_row_is_a_singleton_per_bucket(self, pg_conn): + repo = _repo(pg_conn) + repo.upsert(scope="instance", subject_id=None, token_limit=1) + repo.upsert(scope="instance", subject_id=None, token_limit=2) + repo.upsert(scope="instance", subject_id=None, bucket="agent", token_limit=3) + rows = repo.list_by_scope("instance") + assert [(r["bucket"], r["token_limit"]) for r in rows] == [("all", 2), ("agent", 3)] + + @pytest.mark.parametrize( + "kwargs", + [ + {"scope": "org", "subject_id": "x"}, + {"scope": "user", "subject_id": None}, + {"scope": "instance", "subject_id": "x"}, + {"scope": "user", "subject_id": "u", "bucket": "nope"}, + {"scope": "user", "subject_id": "u", "token_limit": 1, "token_unlimited": True}, + {"scope": "user", "subject_id": "u", "cost_limit_usd": 1, "cost_unlimited": True}, + {"scope": "user", "subject_id": "u", "token_limit": -1}, + {"scope": "user", "subject_id": "u", "cost_limit_usd": -0.5}, + ], + ) + def test_rejects_invalid_policies(self, pg_conn, kwargs): + with pytest.raises(ValueError): + _repo(pg_conn).upsert(**kwargs) + + +class TestReads: + def test_get_and_list_for_subject(self, pg_conn): + repo = _repo(pg_conn) + repo.upsert(scope="user", subject_id="u1", bucket="agent", token_limit=5) + repo.upsert(scope="user", subject_id="u1", token_limit=9) + assert repo.get("user", "u1")["token_limit"] == 9 + assert repo.get("user", "u1", "direct") is None + assert [r["bucket"] for r in repo.list_for_subject("user", "u1")] == ["all", "agent"] + assert repo.list_for_subject("user", "nobody") == [] + + def test_list_by_scope_rejects_unknown_scope(self, pg_conn): + with pytest.raises(ValueError): + _repo(pg_conn).list_by_scope("org") + + +class TestPoliciesForUser: + def test_collects_instance_user_and_team_rows(self, pg_conn): + repo = _repo(pg_conn) + mine, other = _team(pg_conn, "qp-mine"), _team(pg_conn, "qp-other") + _member(pg_conn, mine, "u1") + repo.upsert(scope="instance", subject_id=None, token_limit=1) + repo.upsert(scope="team", subject_id=mine, token_limit=2) + repo.upsert(scope="team", subject_id=other, token_limit=3) + repo.upsert(scope="user", subject_id="u1", token_limit=4) + repo.upsert(scope="user", subject_id="u2", token_limit=5) + + limits = sorted(r["token_limit"] for r in repo.policies_for_user("u1")) + assert limits == [1, 2, 4] + + def test_a_team_counts_once_however_many_memberships(self, pg_conn): + repo = _repo(pg_conn) + team = _team(pg_conn, "qp-multi") + _member(pg_conn, team, "u1", "team_member", "manual") + _member(pg_conn, team, "u1", "team_admin", "manual") + _member(pg_conn, team, "u1", "team_member", "oidc_group") + repo.upsert(scope="team", subject_id=team, token_limit=7) + assert [r["token_limit"] for r in repo.policies_for_user("u1")] == [7] + + def test_every_team_of_the_user_is_included(self, pg_conn): + repo = _repo(pg_conn) + for slug, limit in (("qp-a", 10), ("qp-b", 20), ("qp-c", 30)): + team = _team(pg_conn, slug) + _member(pg_conn, team, "u1") + repo.upsert(scope="team", subject_id=team, token_limit=limit) + assert sorted(r["token_limit"] for r in repo.policies_for_user("u1")) == [10, 20, 30] + + def test_disabled_rows_are_left_out(self, pg_conn): + repo = _repo(pg_conn) + repo.upsert(scope="user", subject_id="u1", token_limit=4, enabled=False) + assert repo.policies_for_user("u1") == [] + + def test_leaving_a_team_drops_its_allowance(self, pg_conn): + repo = _repo(pg_conn) + team = _team(pg_conn, "qp-leave") + _member(pg_conn, team, "u1") + repo.upsert(scope="team", subject_id=team, token_limit=7) + pg_conn.execute(text("DELETE FROM team_members WHERE user_id = 'u1'")) + assert repo.policies_for_user("u1") == [] + + +class TestDelete: + def test_delete_one_bucket_or_all(self, pg_conn): + repo = _repo(pg_conn) + repo.upsert(scope="user", subject_id="u1", token_limit=1) + repo.upsert(scope="user", subject_id="u1", bucket="agent", token_limit=2) + repo.upsert(scope="user", subject_id="u2", token_limit=3) + assert repo.delete("user", "u1", "agent") == 1 + assert repo.delete("user", "u1", "agent") == 0 + assert repo.delete("user", "u1") == 1 + assert repo.get("user", "u2") is not None From 5406550bca8d25550b8ac00664dc3bbe6a5925a4 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 11:51:23 +0100 Subject: [PATCH 05/14] 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 | 10 ++- docsgpt/api/answer/routes/answer.py | 8 +- docsgpt/api/answer/routes/base.py | 18 +++- docsgpt/api/answer/routes/stream.py | 8 +- docsgpt/api/user/scheduler_worker.py | 6 ++ docsgpt/api/v1/routes.py | 2 +- docsgpt/quotas/__init__.py | 3 +- docsgpt/quotas/http.py | 14 +++ docsgpt/quotas/service.py | 8 ++ docsgpt/worker.py | 7 ++ tests/api/user/test_scheduler_worker.py | 25 ++++++ tests/quotas/test_enforcement.py | 108 ++++++++++++++++++++++++ tests/worker/test_agent_workers.py | 30 +++++++ 13 files changed, 238 insertions(+), 9 deletions(-) create mode 100644 docsgpt/quotas/http.py create mode 100644 tests/quotas/test_enforcement.py diff --git a/docsgpt/agents/headless_runner.py b/docsgpt/agents/headless_runner.py index 59c40c5c..344e5fe3 100644 --- a/docsgpt/agents/headless_runner.py +++ b/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") diff --git a/docsgpt/api/answer/routes/answer.py b/docsgpt/api/answer/routes/answer.py index 17fac152..cdc63cca 100644 --- a/docsgpt/api/answer/routes/answer.py +++ b/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( diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index 748f0dda..c61cfd48 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/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: diff --git a/docsgpt/api/answer/routes/stream.py b/docsgpt/api/answer/routes/stream.py index f79e0633..fdc5a00b 100644 --- a/docsgpt/api/answer/routes/stream.py +++ b/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"), diff --git a/docsgpt/api/user/scheduler_worker.py b/docsgpt/api/user/scheduler_worker.py index 6207832d..3e085f0e 100644 --- a/docsgpt/api/user/scheduler_worker.py +++ b/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" diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index 4137d6dc..cd208e84 100644 --- a/docsgpt/api/v1/routes.py +++ b/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 diff --git a/docsgpt/quotas/__init__.py b/docsgpt/quotas/__init__.py index ae1b5f4f..9ed71a75 100644 --- a/docsgpt/quotas/__init__.py +++ b/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", diff --git a/docsgpt/quotas/http.py b/docsgpt/quotas/http.py new file mode 100644 index 00000000..494b87cf --- /dev/null +++ b/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 diff --git a/docsgpt/quotas/service.py b/docsgpt/quotas/service.py index 775f3109..7659fe1d 100644 --- a/docsgpt/quotas/service.py +++ b/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.""" diff --git a/docsgpt/worker.py b/docsgpt/worker.py index 3ce7f15e..947f59c2 100755 --- a/docsgpt/worker.py +++ b/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 diff --git a/tests/api/user/test_scheduler_worker.py b/tests/api/user/test_scheduler_worker.py index 5b7b55ea..09528cd5 100644 --- a/tests/api/user/test_scheduler_worker.py +++ b/tests/api/user/test_scheduler_worker.py @@ -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, ): diff --git a/tests/quotas/test_enforcement.py b/tests/quotas/test_enforcement.py new file mode 100644 index 00000000..b02dbae1 --- /dev/null +++ b/tests/quotas/test_enforcement.py @@ -0,0 +1,108 @@ +"""Tests for quota enforcement at the request entry points.""" + +from __future__ import annotations + +import json +from contextlib import contextmanager +from unittest.mock import patch + +import pytest + +from docsgpt.quotas.service import QuotaExceededError +from docsgpt.storage.db.repositories.agents import AgentsRepository +from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository + + +@pytest.fixture +def db(pg_conn): + @contextmanager + def _yield(): + yield pg_conn + + with patch("docsgpt.quotas.service.db_readonly", _yield), patch( + "docsgpt.api.answer.routes.base.db_readonly", _yield + ), patch("docsgpt.agents.headless_runner.db_readonly", _yield): + yield pg_conn + + +def _spend(conn, user_id, tokens, api_key=None): + TokenUsageRepository(conn).insert(user_id=user_id, api_key=api_key, prompt_tokens=tokens) + + +def _check(flask_app, agent_config, decoded_token=None): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + return BaseAnswerResource().check_usage(agent_config, decoded_token) + + +class TestCheckUsage: + def test_direct_chat_is_limited(self, db, flask_app): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100) + _spend(db, "u1", 100) + + response = _check(flask_app, {}, {"sub": "u1"}) + + assert response.status_code == 429 + assert int(response.headers["Retry-After"]) >= 1 + body = json.loads(response.data) + assert body["error_code"] == "quota-exceeded" + assert (body["dimension"], body["usage"], body["limit"], body["source"]) == ("tokens", 100, 100, "user") + + def test_direct_chat_under_the_limit_passes(self, db, flask_app): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100) + _spend(db, "u1", 99) + assert _check(flask_app, {}, {"sub": "u1"}) is None + + def test_no_policies_and_no_identity_pass(self, db, flask_app): + assert _check(flask_app, {}) is None + assert _check(flask_app, {}, {"sub": "u1"}) is None + + def test_agent_traffic_bills_the_resolved_user(self, db, flask_app): + AgentsRepository(db).create("owner", "a", "published", key="k1") + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="owner", token_limit=10) + _spend(db, "owner", 10, api_key="k1") + + config = {"user_api_key": "k1", "user_id": "owner"} + assert _check(flask_app, config, {"sub": "owner"}).status_code == 429 + # A shared agent bills the caller, who has room. + assert _check(flask_app, config, {"sub": "caller"}) is None + # No token resolved: fall back to the agent owner. + assert _check(flask_app, config).status_code == 429 + + def test_quota_runs_before_the_agent_key_lookup(self, db, flask_app): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=0) + response = _check(flask_app, {"user_api_key": "missing"}, {"sub": "u1"}) + assert response.status_code == 429 + + def test_agent_limits_still_apply(self, db, flask_app): + AgentsRepository(db).create( + "owner", "a", "published", key="k2", limited_token_mode=True, token_limit=5 + ) + _spend(db, "owner", 5, api_key="k2") + response = _check(flask_app, {"user_api_key": "k2"}, {"sub": "owner"}) + assert response.status_code == 429 + assert "error_code" not in json.loads(response.data) + + def test_agent_bucket_policy_ignores_direct_chat(self, db, flask_app): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", bucket="agent", token_limit=10) + _spend(db, "u1", 500) + assert _check(flask_app, {}, {"sub": "u1"}) is None + + +class TestHeadless: + def test_exhausted_owner_is_refused_before_the_run(self, db): + from docsgpt.agents.headless_runner import run_agent_headless + + QuotaPoliciesRepository(db).upsert(scope="instance", subject_id=None, token_limit=10) + _spend(db, "owner", 10) + + with patch("docsgpt.agents.headless_runner.RetrieverCreator") as retriever: + with pytest.raises(QuotaExceededError) as raised: + run_agent_headless({"user_id": "owner", "key": "k"}, "hello") + + retriever.create_retriever.assert_not_called() + assert raised.value.exceeded.source == "instance" + assert "10 of 10 tokens" in str(raised.value) + diff --git a/tests/worker/test_agent_workers.py b/tests/worker/test_agent_workers.py index 8a863f64..033f97d0 100644 --- a/tests/worker/test_agent_workers.py +++ b/tests/worker/test_agent_workers.py @@ -110,6 +110,36 @@ class TestAgentWebhookWorker: with pytest.raises(RuntimeError, match="LLM exploded"): worker.agent_webhook_worker(task_self, agent_id, {"event": "ping"}) + def test_quota_refusal_is_returned_not_retried( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + """A spent quota cannot succeed on retry, so the task must not raise.""" + from datetime import datetime, timezone + + from docsgpt import worker + from docsgpt.agents import headless_runner + from docsgpt.quotas.service import QuotaExceeded, QuotaExceededError + + agent = AgentsRepository(pg_conn).create( + user_id="alice", name="hook-agent", status="active", + agent_type="classic", retriever="classic", chunks=2, key="sk-test-q", + ) + exceeded = QuotaExceeded( + user_id="alice", bucket="all", budget="cost", usage=5.0, limit=5.0, + source="user", source_id=None, + resets_at=datetime(2099, 1, 1, tzinfo=timezone.utc), + ) + + def _refuse(*a, **k): + raise QuotaExceededError(exceeded) + + monkeypatch.setattr(headless_runner, "run_agent_headless", _refuse) + + result = worker.agent_webhook_worker(task_self, str(agent["id"]), {"event": "ping"}) + + assert result["status"] == "quota_exceeded" + assert "$5.00 of $5.00" in result["error"] + def test_webhook_journals_headless_denial_for_approval_gated_tool( self, pg_conn, patch_worker_db, task_self, monkeypatch ): From 21f43b2966d93bbc906082673c9e4380cda52c87 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 11:54:44 +0100 Subject: [PATCH 06/14] feat(quotas): admin quota API and GET /api/user/quota Admins read and set the instance default, team allowances and user overrides under /api/admin/quotas. A user's endpoint also returns the limits those layers resolve to, the layer each came from and the usage against them. The overview lists catalog models used this period that have no price, since a cost limit cannot see them. Every write is audited. GET /api/user/quota gives a user their own limited buckets, usage and reset time without naming the policies behind them; any valid token may call it. --- docsgpt/api/admin/__init__.py | 1 + docsgpt/api/admin/quotas.py | 276 +++++++++++++++++ docsgpt/api/pat/rules.py | 1 + docsgpt/api/user/me/routes.py | 36 +++ .../storage/db/repositories/token_usage.py | 22 ++ tests/api/test_quota_endpoints.py | 280 ++++++++++++++++++ 6 files changed, 616 insertions(+) create mode 100644 docsgpt/api/admin/quotas.py create mode 100644 tests/api/test_quota_endpoints.py diff --git a/docsgpt/api/admin/__init__.py b/docsgpt/api/admin/__init__.py index 8fa8b444..7bfc3059 100644 --- a/docsgpt/api/admin/__init__.py +++ b/docsgpt/api/admin/__init__.py @@ -1,3 +1,4 @@ from .routes import admin_ns +from . import quotas # noqa: F401 (registers the quota resources on admin_ns) __all__ = ["admin_ns"] diff --git a/docsgpt/api/admin/quotas.py b/docsgpt/api/admin/quotas.py new file mode 100644 index 00000000..8210ba8d --- /dev/null +++ b/docsgpt/api/admin/quotas.py @@ -0,0 +1,276 @@ +"""Admin endpoints for usage quotas (RBAC ``admin`` role required). + +Policies are set at three layers: the instance default, a team's per-member +allowance and a single user's override. Every write is audited to +``auth_events`` with the acting admin recorded. +""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +from flask import jsonify, make_response, request +from flask_restx import Resource + +from docsgpt.api.admin.routes import _actor, admin_ns +from docsgpt.api.user.authz import admin_required +from docsgpt.core.settings import settings +from docsgpt.pricing import is_priced +from docsgpt.quotas.service import REQUEST_BUCKETS, QuotaService +from docsgpt.quotas.windows import window_bounds +from docsgpt.storage.db.base_repository import looks_like_uuid +from docsgpt.storage.db.repositories.auth_events import AuthEventsRepository +from docsgpt.storage.db.repositories.quota_policies import BUCKETS, QuotaPoliciesRepository +from docsgpt.storage.db.repositories.teams import TeamsRepository +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository +from docsgpt.storage.db.repositories.users import UsersRepository +from docsgpt.storage.db.session import db_readonly, db_session + +_MAX_TOKEN_LIMIT = 2**62 +_MAX_COST_LIMIT = 99_999_999.0 +_MAX_NOTE_LENGTH = 500 + + +class _BadPolicy(ValueError): + """The request body does not describe a valid policy.""" + + +def _policy_json(row: dict) -> dict: + cost = row.get("cost_limit_usd") + return { + "scope": row["scope"], + "subject_id": row.get("subject_id"), + "bucket": row["bucket"], + "token_limit": row.get("token_limit"), + "token_unlimited": bool(row.get("token_unlimited")), + "cost_limit_usd": float(cost) if cost is not None else None, + "cost_unlimited": bool(row.get("cost_unlimited")), + "enabled": bool(row.get("enabled", True)), + "note": row.get("note"), + "updated_by": row.get("updated_by"), + "updated_at": row.get("updated_at"), + } + + +def _error(message: str, status: int): + return make_response(jsonify({"success": False, "message": message}), status) + + +def _bucket(value: Any) -> str: + if value is None: + return "all" + if value not in BUCKETS: + raise _BadPolicy(f"bucket must be one of: {', '.join(BUCKETS)}") + return value + + +def _flag(data: dict, key: str, default: bool) -> bool: + value = data.get(key, default) + if not isinstance(value, bool): + raise _BadPolicy(f"{key} must be a boolean") + return value + + +def _parse_policy(data: Any) -> dict: + """Validate a policy body into ``QuotaPoliciesRepository.upsert`` kwargs.""" + if not isinstance(data, dict): + raise _BadPolicy("Body must be a JSON object") + token_limit = data.get("token_limit") + if token_limit is not None: + if isinstance(token_limit, bool) or not isinstance(token_limit, int): + raise _BadPolicy("token_limit must be a whole number or null") + if not 0 <= token_limit <= _MAX_TOKEN_LIMIT: + raise _BadPolicy("token_limit is out of range") + cost_limit = data.get("cost_limit_usd") + if cost_limit is not None: + if isinstance(cost_limit, bool) or not isinstance(cost_limit, (int, float)): + raise _BadPolicy("cost_limit_usd must be a number or null") + if not math.isfinite(cost_limit) or not 0 <= cost_limit <= _MAX_COST_LIMIT: + raise _BadPolicy("cost_limit_usd is out of range") + cost_limit = round(float(cost_limit), 4) + token_unlimited = _flag(data, "token_unlimited", False) + cost_unlimited = _flag(data, "cost_unlimited", False) + if token_unlimited and token_limit is not None: + raise _BadPolicy("Set token_limit or token_unlimited, not both") + if cost_unlimited and cost_limit is not None: + raise _BadPolicy("Set cost_limit_usd or cost_unlimited, not both") + if token_limit is None and cost_limit is None and not token_unlimited and not cost_unlimited: + raise _BadPolicy("Set a limit or mark a budget unlimited; delete the policy to remove it") + note = data.get("note") + if note is not None: + if not isinstance(note, str): + raise _BadPolicy("note must be a string") + note = note.strip()[:_MAX_NOTE_LENGTH] or None + return { + "bucket": _bucket(data.get("bucket")), + "token_limit": token_limit, + "token_unlimited": token_unlimited, + "cost_limit_usd": cost_limit, + "cost_unlimited": cost_unlimited, + "enabled": _flag(data, "enabled", True), + "note": note, + } + + +def _audit(conn, event: str, scope: str, subject_id: Optional[str], detail: dict) -> None: + actor = _actor() + AuthEventsRepository(conn).insert( + # A user policy is filed under that user; the rest under the acting admin. + subject_id if scope == "user" else (actor or "unknown"), + event, + ip=request.remote_addr, + user_agent=request.headers.get("User-Agent"), + metadata={"by": actor, "via": "admin_api", "scope": scope, "subject_id": subject_id, **detail}, + ) + + +def _put_policy(scope: str, subject_id: Optional[str]): + try: + fields = _parse_policy(request.get_json(silent=True)) + except _BadPolicy as exc: + return _error(str(exc), 400) + with db_session() as conn: + row = QuotaPoliciesRepository(conn).upsert( + scope=scope, subject_id=subject_id, actor=_actor(), **fields + ) + _audit(conn, "quota_policy_set", scope, subject_id, fields) + return make_response(jsonify({"success": True, "policy": _policy_json(row)}), 200) + + +def _delete_policy(scope: str, subject_id: Optional[str]): + try: + bucket = _bucket(request.args.get("bucket")) if "bucket" in request.args else None + except _BadPolicy as exc: + return _error(str(exc), 400) + with db_session() as conn: + deleted = QuotaPoliciesRepository(conn).delete(scope, subject_id, bucket) + if deleted: + _audit(conn, "quota_policy_deleted", scope, subject_id, {"bucket": bucket or "*"}) + return make_response(jsonify({"success": True, "deleted": deleted}), 200) + + +def _unpriced_models(conn) -> list[dict]: + """Catalog models used this period that no cost limit can see.""" + start, _ = window_bounds(settings.QUOTA_PERIOD) + return [ + row + for row in TokenUsageRepository(conn).tokens_by_model(start=start) + # BYOM ids are UUIDs; those calls are $0 by design, not by omission. + if not looks_like_uuid(row["model_id"]) and not is_priced(row["model_id"]) + ] + + +@admin_ns.route("/admin/quotas") +class AdminQuotasResource(Resource): + @admin_required + def get(self): + """Every stored policy, grouped by layer, plus the models cost limits cannot see.""" + start, resets_at = window_bounds(settings.QUOTA_PERIOD) + with db_readonly() as conn: + repo = QuotaPoliciesRepository(conn) + teams = {str(t["id"]): t for t in TeamsRepository(conn).list_all()} + team_policies = [] + for row in repo.list_by_scope("team"): + team = teams.get(str(row["subject_id"]), {}) + team_policies.append( + { + **_policy_json(row), + "team_name": team.get("name"), + "team_slug": team.get("slug"), + "member_count": team.get("member_count"), + } + ) + body = { + "success": True, + "period": settings.QUOTA_PERIOD, + "period_start": start.isoformat(), + "resets_at": resets_at.isoformat(), + "instance": [_policy_json(r) for r in repo.list_by_scope("instance")], + "teams": team_policies, + "users": [_policy_json(r) for r in repo.list_by_scope("user")], + "unpriced_models": _unpriced_models(conn), + } + return make_response(jsonify(body), 200) + + +@admin_ns.route("/admin/quotas/instance") +class AdminInstanceQuotaResource(Resource): + @admin_required + def put(self): + """Set the instance default for one bucket.""" + return _put_policy("instance", None) + + @admin_required + def delete(self): + """Remove the instance default for ``?bucket=``, or for every bucket.""" + return _delete_policy("instance", None) + + +@admin_ns.route("/admin/quotas/teams/") +class AdminTeamQuotaResource(Resource): + @admin_required + def get(self, team_id): + """The per-member allowance of one team.""" + if not looks_like_uuid(team_id): + return _error("Team not found", 404) + with db_readonly() as conn: + if TeamsRepository(conn).get(team_id) is None: + return _error("Team not found", 404) + rows = QuotaPoliciesRepository(conn).list_for_subject("team", team_id) + return make_response( + jsonify({"success": True, "policies": [_policy_json(r) for r in rows]}), 200 + ) + + @admin_required + def put(self, team_id): + """Set the allowance each member of the team gets.""" + if not looks_like_uuid(team_id): + return _error("Team not found", 404) + with db_readonly() as conn: + if TeamsRepository(conn).get(team_id) is None: + return _error("Team not found", 404) + return _put_policy("team", team_id) + + @admin_required + def delete(self, team_id): + """Remove the team's allowance for ``?bucket=``, or for every bucket.""" + if not looks_like_uuid(team_id): + return _error("Team not found", 404) + return _delete_policy("team", team_id) + + +@admin_ns.route("/admin/quotas/users/") +class AdminUserQuotaResource(Resource): + @admin_required + def get(self, user_id): + """A user's overrides and the limits and usage they resolve to.""" + with db_readonly() as conn: + if UsersRepository(conn).get(user_id) is None: + return _error("User not found", 404) + rows = QuotaPoliciesRepository(conn).list_for_subject("user", user_id) + statuses = QuotaService.status(user_id, ("all", *REQUEST_BUCKETS)) + return make_response( + jsonify( + { + "success": True, + "period": settings.QUOTA_PERIOD, + "policies": [_policy_json(r) for r in rows], + "effective": [s.to_dict() for s in statuses], + } + ), + 200, + ) + + @admin_required + def put(self, user_id): + """Set one user's override, which beats team allowances and the default.""" + with db_readonly() as conn: + if UsersRepository(conn).get(user_id) is None: + return _error("User not found", 404) + return _put_policy("user", user_id) + + @admin_required + def delete(self, user_id): + """Remove the user's override for ``?bucket=``, or for every bucket.""" + return _delete_policy("user", user_id) diff --git a/docsgpt/api/pat/rules.py b/docsgpt/api/pat/rules.py index 9b69d1a7..227fe11a 100644 --- a/docsgpt/api/pat/rules.py +++ b/docsgpt/api/pat/rules.py @@ -161,6 +161,7 @@ _CHAT = dict( RULES: dict[tuple[str, str], Rule] = { # Identity and public metadata: any valid token. ("/api/user/me", "GET"): _rule(open=True), + ("/api/user/quota", "GET"): _rule(open=True), ("/api/health", "GET"): _rule(open=True), ("/api/config", "GET"): _rule(open=True), # Agents diff --git a/docsgpt/api/user/me/routes.py b/docsgpt/api/user/me/routes.py index 67e31dfb..a6c70d4b 100644 --- a/docsgpt/api/user/me/routes.py +++ b/docsgpt/api/user/me/routes.py @@ -5,6 +5,9 @@ only from ``request.decoded_token`` (already populated and role-resolved by the auth chokepoint in ``app.py``). Auth-mode-agnostic. ``email``/``name``/ ``picture`` are OIDC-only and optional — they are echoed from the token and are never present for ``simple_jwt``/``session_jwt``/no-auth modes. + +``GET /api/user/quota`` returns the caller's usage against the limits an admin +set for them, without naming the policies behind those limits. """ from __future__ import annotations @@ -13,6 +16,8 @@ from flask import jsonify, make_response, request from flask_restx import Namespace, Resource from docsgpt.api.pat.tokens import is_pat +from docsgpt.core.settings import settings +from docsgpt.quotas.service import REQUEST_BUCKETS, QuotaService me_ns = Namespace("me", description="Current user identity and roles", path="/api") @@ -43,3 +48,34 @@ class MeResource(Resource): "resource_filter": decoded_token.get("resource_filter") or {}, } return make_response(jsonify(body), 200) + + +def _own_budget(budget: dict) -> dict: + return {"limit": budget["limit"], "used": budget["used"]} + + +@me_ns.route("/user/quota") +class MyQuotaResource(Resource): + def get(self): + """Return the caller's limited buckets: ``{bucket, tokens, cost, resets_at}`` each.""" + decoded_token = getattr(request, "decoded_token", None) + user_id = decoded_token.get("sub") if decoded_token else None + if not user_id: + return make_response(jsonify({"success": False}), 401) + statuses = QuotaService.status(user_id, ("all", *REQUEST_BUCKETS)) + buckets = [] + for status in statuses: + if status.limits.unlimited: + continue + data = status.to_dict() + buckets.append( + { + "bucket": data["bucket"], + "tokens": _own_budget(data["tokens"]), + "cost": _own_budget(data["cost"]), + "resets_at": data["resets_at"], + } + ) + return make_response( + jsonify({"success": True, "period": settings.QUOTA_PERIOD, "buckets": buckets}), 200 + ) diff --git a/docsgpt/storage/db/repositories/token_usage.py b/docsgpt/storage/db/repositories/token_usage.py index 17e546fe..2c8038fa 100644 --- a/docsgpt/storage/db/repositories/token_usage.py +++ b/docsgpt/storage/db/repositories/token_usage.py @@ -165,6 +165,28 @@ class TokenUsageRepository: ).one() return int(row[0]), float(row[1]) + def tokens_by_model(self, *, start: datetime) -> list[dict]: + """Return ``{model_id, tokens, cost}`` per model since ``start``, busiest first.""" + result = self._conn.execute( + text( + """ + SELECT model_id, + COALESCE(SUM(prompt_tokens + generated_tokens), 0) AS tokens, + COALESCE(SUM(cost), 0) AS cost + FROM token_usage + WHERE timestamp >= :start AND model_id IS NOT NULL + AND source <> ALL(:rollup_sources) + GROUP BY model_id + ORDER BY tokens DESC, model_id + """ + ), + {"start": start, "rollup_sources": list(self.ROLLUP_SOURCES)}, + ) + return [ + {"model_id": row[0], "tokens": int(row[1]), "cost": float(row[2])} + for row in result.fetchall() + ] + # Token usage written outside a user-initiated request (conversation # title generation, history compression, RAG question condensing, # provider fallback). Mirrors the exclusion list in ``count_in_range``. diff --git a/tests/api/test_quota_endpoints.py b/tests/api/test_quota_endpoints.py new file mode 100644 index 00000000..44c424db --- /dev/null +++ b/tests/api/test_quota_endpoints.py @@ -0,0 +1,280 @@ +"""Endpoint tests for the admin quota API and ``GET /api/user/quota``. + +Driven through the real app.py chokepoint against an ephemeral Postgres; only +``handle_auth`` / ``resolve_roles`` are patched. +""" + +from __future__ import annotations + +import json +from contextlib import ExitStack, contextmanager +from unittest.mock import patch + +import pytest +from sqlalchemy import text + +from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository +from docsgpt.storage.db.repositories.users import UsersRepository + + +@pytest.fixture +def client(): + from docsgpt.app import app as flask_app + + flask_app.config["TESTING"] = True + return flask_app.test_client() + + +@pytest.fixture +def db(pg_conn): + @contextmanager + def _yield(): + yield pg_conn + + with ExitStack() as stack: + for target in ( + "docsgpt.api.admin.quotas.db_readonly", + "docsgpt.api.admin.quotas.db_session", + "docsgpt.quotas.service.db_readonly", + ): + stack.enter_context(patch(target, _yield)) + yield pg_conn + + +@contextmanager +def _as(sub, *roles): + with patch("docsgpt.app.handle_auth", return_value={"sub": sub}), patch( + "docsgpt.app.resolve_roles", return_value=list(roles) or ["user"] + ): + yield + + +def _admin(): + return _as("admin1", "admin", "user") + + +def _body(resp): + return json.loads(resp.data) + + +def _team(conn, slug="q-team", member=None): + team_id = str( + conn.execute( + text("INSERT INTO teams (name, slug, owner_id) VALUES (:n, :s, 'o') RETURNING id"), + {"n": slug, "s": slug}, + ).scalar() + ) + if member: + conn.execute( + text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), :u, 'team_member')"), + {"t": team_id, "u": member}, + ) + return team_id + + +ADMIN_ROUTES = [ + ("get", "/api/admin/quotas"), + ("put", "/api/admin/quotas/instance"), + ("delete", "/api/admin/quotas/instance"), + ("get", "/api/admin/quotas/users/u1"), + ("put", "/api/admin/quotas/users/u1"), + ("delete", "/api/admin/quotas/users/u1"), + ("get", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"), + ("put", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"), + ("delete", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"), +] + + +class TestGuard: + @pytest.mark.parametrize("method, path", ADMIN_ROUTES) + def test_non_admin_forbidden(self, client, db, method, path): + with _as("u1"): + assert getattr(client, method)(path, json={"token_limit": 1}).status_code == 403 + + @pytest.mark.parametrize("method, path", ADMIN_ROUTES) + def test_unauthenticated(self, client, method, path): + with patch("docsgpt.app.handle_auth", return_value=None): + assert getattr(client, method)(path, json={"token_limit": 1}).status_code == 401 + + def test_team_admin_cannot_set_their_teams_allowance(self, client, db): + team_id = _team(db) + db.execute( + text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), 'lead', 'team_admin')"), + {"t": team_id}, + ) + with _as("lead"): + resp = client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_unlimited": True}) + assert resp.status_code == 403 + assert QuotaPoliciesRepository(db).get("team", team_id) is None + + +class TestInstancePolicy: + def test_set_read_delete(self, client, db): + with _admin(): + put = client.put("/api/admin/quotas/instance", json={"token_limit": 1000, "note": " default "}) + assert put.status_code == 200 + policy = _body(put)["policy"] + assert (policy["token_limit"], policy["note"], policy["updated_by"]) == (1000, "default", "admin1") + + client.put("/api/admin/quotas/instance", json={"bucket": "agent", "cost_limit_usd": 2.5}) + overview = _body(client.get("/api/admin/quotas")) + assert [(p["bucket"], p["token_limit"], p["cost_limit_usd"]) for p in overview["instance"]] == [ + ("all", 1000, None), + ("agent", None, 2.5), + ] + assert overview["period"] == "month" + + assert _body(client.delete("/api/admin/quotas/instance?bucket=agent"))["deleted"] == 1 + assert _body(client.delete("/api/admin/quotas/instance"))["deleted"] == 1 + assert _body(client.get("/api/admin/quotas"))["instance"] == [] + + def test_writes_are_audited(self, client, db): + with _admin(): + client.put("/api/admin/quotas/instance", json={"token_limit": 5}) + client.delete("/api/admin/quotas/instance") + client.delete("/api/admin/quotas/instance") + events = db.execute( + text("SELECT user_id, event, metadata FROM auth_events WHERE event LIKE 'quota_policy_%'") + ).fetchall() + by_event = {e[1]: e for e in events} + # The second delete removed nothing, so it left no event. + assert sorted((e[0], e[1]) for e in events) == [ + ("admin1", "quota_policy_deleted"), + ("admin1", "quota_policy_set"), + ] + metadata = by_event["quota_policy_set"][2] + assert metadata["token_limit"] == 5 and metadata["by"] == "admin1" + + @pytest.mark.parametrize( + "body", + [ + None, + [], + {}, + {"note": "only a note"}, + {"token_limit": -1}, + {"token_limit": 1.5}, + {"token_limit": True}, + {"token_limit": "10"}, + {"token_limit": 2**63}, + {"cost_limit_usd": -0.01}, + {"cost_limit_usd": "5"}, + {"cost_limit_usd": float("inf")}, + {"cost_limit_usd": 1e12}, + {"token_limit": 1, "token_unlimited": True}, + {"cost_limit_usd": 1, "cost_unlimited": True}, + {"token_unlimited": "yes"}, + {"token_limit": 1, "enabled": "no"}, + {"token_limit": 1, "bucket": "everything"}, + {"token_limit": 1, "note": 7}, + ], + ) + def test_invalid_bodies_rejected(self, client, db, body): + with _admin(): + resp = client.put("/api/admin/quotas/instance", json=body) + assert resp.status_code == 400 + assert QuotaPoliciesRepository(db).list_by_scope("instance") == [] + + def test_unknown_bucket_on_delete(self, client, db): + with _admin(): + assert client.delete("/api/admin/quotas/instance?bucket=nope").status_code == 400 + + def test_zero_is_accepted_as_a_block(self, client, db): + with _admin(): + resp = client.put("/api/admin/quotas/instance", json={"token_limit": 0, "cost_limit_usd": 0}) + assert resp.status_code == 200 + assert _body(resp)["policy"]["token_limit"] == 0 + + +class TestTeamPolicy: + def test_set_and_list_with_team_details(self, client, db): + team_id = _team(db, "q-eng", member="u1") + with _admin(): + assert client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_limit": 500}).status_code == 200 + (row,) = _body(client.get("/api/admin/quotas"))["teams"] + assert (row["team_slug"], row["token_limit"], row["member_count"]) == ("q-eng", 500, 1) + assert _body(client.get(f"/api/admin/quotas/teams/{team_id}"))["policies"][0]["token_limit"] == 500 + + @pytest.mark.parametrize("team_id", ["not-a-uuid", "00000000-0000-0000-0000-000000000000"]) + def test_unknown_team(self, client, db, team_id): + with _admin(): + assert client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_limit": 1}).status_code == 404 + assert client.get(f"/api/admin/quotas/teams/{team_id}").status_code == 404 + + +class TestUserPolicy: + def test_unknown_user(self, client, db): + with _admin(): + assert client.put("/api/admin/quotas/users/ghost", json={"token_limit": 1}).status_code == 404 + assert client.get("/api/admin/quotas/users/ghost").status_code == 404 + + def test_effective_limits_name_their_source(self, client, db): + UsersRepository(db).upsert("u1") + small, big = _team(db, "q-small", member="u1"), _team(db, "q-big", member="u1") + repo = QuotaPoliciesRepository(db) + repo.upsert(scope="instance", subject_id=None, token_limit=10, cost_limit_usd=1) + repo.upsert(scope="team", subject_id=small, token_limit=100) + repo.upsert(scope="team", subject_id=big, token_limit=900) + TokenUsageRepository(db).insert(user_id="u1", prompt_tokens=40, cost=0.25) + + with _admin(): + body = _body(client.get("/api/admin/quotas/users/u1")) + overall = body["effective"][0] + assert overall["bucket"] == "all" + assert overall["tokens"] == {"limit": 900.0, "used": 40, "source": "team", "source_id": big} + assert overall["cost"] == {"limit": 1.0, "used": 0.25, "source": "instance", "source_id": None} + assert body["policies"] == [] + + with _admin(): + client.put("/api/admin/quotas/users/u1", json={"token_limit": 50}) + body = _body(client.get("/api/admin/quotas/users/u1")) + assert body["effective"][0]["tokens"]["source"] == "user" + assert body["policies"][0]["token_limit"] == 50 + + def test_user_policy_audit_is_filed_under_the_user(self, client, db): + UsersRepository(db).upsert("u1") + with _admin(): + client.put("/api/admin/quotas/users/u1", json={"cost_unlimited": True}) + row = db.execute( + text("SELECT user_id, metadata FROM auth_events WHERE event = 'quota_policy_set'") + ).one() + assert row[0] == "u1" and row[1]["by"] == "admin1" and row[1]["scope"] == "user" + + +class TestUnpricedModels: + def test_lists_used_catalog_models_without_a_price(self, client, db): + usage = TokenUsageRepository(db) + usage.insert(user_id="u1", prompt_tokens=10, model_id="local-llama") + usage.insert(user_id="u1", prompt_tokens=5, model_id="claude-haiku-4-5", cost=0.1) + usage.insert(user_id="u1", prompt_tokens=7, model_id="7d0c1a52-2f5e-4c53-9a0e-111111111111") + with _admin(), patch("docsgpt.api.admin.quotas.is_priced", lambda m: m == "claude-haiku-4-5"): + unpriced = _body(client.get("/api/admin/quotas"))["unpriced_models"] + assert unpriced == [{"model_id": "local-llama", "tokens": 10, "cost": 0.0}] + + +class TestMyQuota: + def test_unlimited_user_sees_no_buckets(self, client, db): + with _as("u1"): + body = _body(client.get("/api/user/quota")) + assert body == {"success": True, "period": "month", "buckets": []} + + def test_limited_user_sees_usage_without_policy_internals(self, client, db): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100, note="secret note") + TokenUsageRepository(db).insert(user_id="u1", prompt_tokens=30) + with _as("u1"): + body = _body(client.get("/api/user/quota")) + (bucket,) = body["buckets"] + assert bucket["bucket"] == "all" + assert bucket["tokens"] == {"limit": 100.0, "used": 30} + assert bucket["cost"] == {"limit": None, "used": 0.0} + assert "source" not in json.dumps(body) and "secret" not in json.dumps(body) + + def test_a_user_only_sees_their_own_quota(self, client, db): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100) + with _as("u2"): + assert _body(client.get("/api/user/quota"))["buckets"] == [] + + def test_unauthenticated(self, client): + with patch("docsgpt.app.handle_auth", return_value=None): + assert client.get("/api/user/quota").status_code == 401 From 1c9c94eba72afbd92bc2ff09948f53203d26457a Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 12:01:48 +0100 Subject: [PATCH 07/14] feat(frontend): admin quota management, a usage meter and the quota chat error The admin dashboard gets a Quotas tab for the instance default, team allowances and user overrides, with a notice listing models that cost limits cannot see. A Quota action on the Users tab shows a user's effective limits, the layer each comes from and their usage, next to the editor for their override. Each budget is either not set at that layer, a limit, or unlimited. Users with a quota see a usage meter with the reset time on the Analytics settings page, and a refused chat request shows the used amount, the limit and the reset time in the user's language. --- frontend/src/admin/AdminUI.tsx | 2 + frontend/src/admin/QuotaEditor.tsx | 256 +++++++++++++ frontend/src/admin/Quotas.tsx | 343 ++++++++++++++++++ frontend/src/admin/UserQuotaModal.tsx | 115 ++++++ frontend/src/admin/Users.tsx | 17 +- frontend/src/admin/index.tsx | 3 + frontend/src/admin/quotaUtils.test.ts | 120 ++++++ frontend/src/admin/quotaUtils.ts | 130 +++++++ frontend/src/api/endpoints.ts | 7 + frontend/src/api/services/adminService.ts | 25 ++ frontend/src/api/services/userService.ts | 2 + .../src/conversation/conversationHandlers.ts | 7 +- frontend/src/conversation/quotaError.test.ts | 52 +++ frontend/src/conversation/quotaError.ts | 47 +++ frontend/src/locale/de.json | 11 + frontend/src/locale/en.json | 11 + frontend/src/locale/es.json | 11 + frontend/src/locale/jp.json | 11 + frontend/src/locale/ru.json | 11 + frontend/src/locale/zh-TW.json | 11 + frontend/src/locale/zh.json | 11 + frontend/src/settings/Analytics.tsx | 4 +- .../src/settings/components/UsageQuota.tsx | 130 +++++++ 23 files changed, 1334 insertions(+), 3 deletions(-) create mode 100644 frontend/src/admin/QuotaEditor.tsx create mode 100644 frontend/src/admin/Quotas.tsx create mode 100644 frontend/src/admin/UserQuotaModal.tsx create mode 100644 frontend/src/admin/quotaUtils.test.ts create mode 100644 frontend/src/admin/quotaUtils.ts create mode 100644 frontend/src/conversation/quotaError.test.ts create mode 100644 frontend/src/conversation/quotaError.ts create mode 100644 frontend/src/settings/components/UsageQuota.tsx diff --git a/frontend/src/admin/AdminUI.tsx b/frontend/src/admin/AdminUI.tsx index e62ab82b..2e5be43e 100644 --- a/frontend/src/admin/AdminUI.tsx +++ b/frontend/src/admin/AdminUI.tsx @@ -119,6 +119,8 @@ const EVENT_LABELS: Record = { scim_created: 'Provisioned', scim_deactivated: 'Deactivated (SCIM)', scim_activated: 'Activated (SCIM)', + quota_policy_set: 'Quota set', + quota_policy_deleted: 'Quota removed', }; export function eventLabel(event: string): string { diff --git a/frontend/src/admin/QuotaEditor.tsx b/frontend/src/admin/QuotaEditor.tsx new file mode 100644 index 00000000..e962ed46 --- /dev/null +++ b/frontend/src/admin/QuotaEditor.tsx @@ -0,0 +1,256 @@ +import { useEffect, useState } from 'react'; +import { useSelector } from 'react-redux'; + +import adminService, { type QuotaScope } from '../api/services/adminService'; +import { Button } from '../components/ui/button'; +import { Input } from '../components/ui/input'; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from '../components/ui/select'; +import { selectToken } from '../preferences/preferenceSlice'; +import { fmtNumber } from './AdminUI'; +import { + fmtUsd, + formToPolicy, + isEmptyForm, + policyToForm, + usagePercent, + type Budget, + type BudgetMode, + type QuotaForm, + type QuotaPolicy, +} from './quotaUtils'; + +const MODES: { value: BudgetMode; label: string }[] = [ + { value: 'inherit', label: 'Not set here' }, + { value: 'limit', label: 'Limit' }, + { value: 'unlimited', label: 'Unlimited' }, +]; + +export function UsageBar({ + label, + budget, + kind, + caption, +}: { + label: string; + budget: Budget; + kind: 'tokens' | 'cost'; + caption?: string; +}) { + const fmt = (n: number) => (kind === 'cost' ? fmtUsd(n) : fmtNumber(n)); + const percent = usagePercent(budget.used, budget.limit); + const tone = + percent >= 100 + ? 'bg-red-500' + : percent >= 80 + ? 'bg-amber-500' + : 'bg-[#7D54D1]'; + return ( +
+
+ {label} + + {fmt(budget.used)} + {budget.limit === null ? ' · no limit' : ` of ${fmt(budget.limit)}`} + +
+ {budget.limit !== null ? ( +
+
+
+ ) : null} + {caption ? ( +

{caption}

+ ) : null} +
+ ); +} + +function BudgetField({ + label, + hint, + mode, + value, + step, + onMode, + onValue, +}: { + label: string; + hint: string; + mode: BudgetMode; + value: string; + step: string; + onMode: (mode: BudgetMode) => void; + onValue: (value: string) => void; +}) { + return ( +
+

{label}

+
+ + {mode === 'limit' ? ( + onValue(e.target.value)} + className="flex-1" + /> + ) : null} +
+
+ ); +} + +/** + * Edits the ``all``-bucket policy of one subject. Saving a form with neither + * budget set removes the policy, since a policy without an opinion is not stored. + */ +export default function QuotaEditor({ + scope, + subjectId, + policy, + inheritHint, + onSaved, +}: { + scope: QuotaScope; + subjectId: string | null; + policy: QuotaPolicy | null; + inheritHint: string; + onSaved: () => void; +}) { + const token = useSelector(selectToken); + const [form, setForm] = useState(() => policyToForm(policy)); + const [busy, setBusy] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + setForm(policyToForm(policy)); + setError(null); + }, [policy, scope, subjectId]); + + const patch = (fields: Partial) => + setForm((prev) => ({ ...prev, ...fields })); + + const submit = async (request: () => Promise) => { + setBusy(true); + setError(null); + try { + const res = await request(); + const json = await res.json().catch(() => ({})); + if (res.ok && json.success !== false) onSaved(); + else setError(json.message || 'Could not save the quota.'); + } catch { + setError('Could not save the quota.'); + } finally { + setBusy(false); + } + }; + + const remove = () => + submit(() => adminService.deleteQuota(scope, subjectId, 'all', token)); + + const save = () => { + if (isEmptyForm(form)) { + if (policy) remove(); + return; + } + const result = formToPolicy(form); + if (!result.ok) { + setError(result.error); + return; + } + submit(() => adminService.setQuota(scope, subjectId, result.policy, token)); + }; + + return ( +
+

{inheritHint}

+ patch({ tokenMode })} + onValue={(tokenLimit) => patch({ tokenLimit })} + /> + patch({ costMode })} + onValue={(costLimit) => patch({ costLimit })} + /> +
+

Note

+ patch({ note: e.target.value })} + className="mt-1" + /> +
+ {error ? ( +

+ {error} +

+ ) : null} +
+ {policy ? ( + + ) : null} + +
+
+ ); +} diff --git a/frontend/src/admin/Quotas.tsx b/frontend/src/admin/Quotas.tsx new file mode 100644 index 00000000..7d5d93ed --- /dev/null +++ b/frontend/src/admin/Quotas.tsx @@ -0,0 +1,343 @@ +import { useCallback, useEffect, useMemo, useState } from 'react'; +import { useSelector } from 'react-redux'; + +import adminService, { type QuotaScope } from '../api/services/adminService'; +import teamsService from '../api/services/teamsService'; +import { Button } from '../components/ui/button'; +import { Modal } from '../components/ui/modal'; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from '../components/ui/select'; +import { + Table, + TableBody, + TableCell, + TableContainer, + TableHead, + TableHeader, + TableRow, +} from '../components/ui/table'; +import { selectToken } from '../preferences/preferenceSlice'; +import { + LoadError, + Loading, + Pill, + fmtDate, + fmtNumber, + fmtRelative, +} from './AdminUI'; +import QuotaEditor from './QuotaEditor'; +import { describeBudget, type QuotaPolicy } from './quotaUtils'; + +type TeamPolicy = QuotaPolicy & { + team_name?: string | null; + team_slug?: string | null; + member_count?: number | null; +}; + +type Editing = { + scope: QuotaScope; + subjectId: string | null; + title: string; + policy: QuotaPolicy | null; +}; + +const HINTS: Record = { + instance: + 'Applies to every user that no team allowance or user override covers.', + team: 'Each member gets this allowance; it is not a shared pool. A member of several teams gets the most generous one.', + user: 'Overrides team allowances and the instance default for this user.', +}; + +function PolicyCells({ policy }: { policy: QuotaPolicy }) { + return ( + <> + + {describeBudget(policy.token_limit, policy.token_unlimited, 'tokens')} + + + {describeBudget(policy.cost_limit_usd, policy.cost_unlimited, 'cost')} + + + {policy.note || '—'} + + + {fmtRelative(policy.updated_at)} + + + ); +} + +export default function Quotas() { + const token = useSelector(selectToken); + const [data, setData] = useState(null); + const [teams, setTeams] = useState([]); + const [loading, setLoading] = useState(true); + const [editing, setEditing] = useState(null); + const [teamPick, setTeamPick] = useState(''); + + const load = useCallback(async () => { + setLoading(true); + try { + const [quotasRes, teamsJson] = await Promise.all([ + adminService.getQuotas(token), + teamsService.listAll(token).catch(() => ({})), + ]); + setData(await quotasRes.json().catch(() => ({ success: false }))); + setTeams(teamsJson?.teams ?? []); + } catch { + setData({ success: false }); + } finally { + setLoading(false); + } + }, [token]); + + useEffect(() => { + load(); + }, [load]); + + // The editor covers the ``all`` bucket; other buckets are listed read-only. + const isAll = (p: QuotaPolicy) => p.bucket === 'all'; + const instancePolicy: QuotaPolicy | null = + (data?.instance ?? []).find(isAll) ?? null; + const teamPolicies: TeamPolicy[] = data?.teams ?? []; + const userPolicies: QuotaPolicy[] = data?.users ?? []; + const teamsWithoutPolicy = useMemo(() => { + const covered = new Set( + teamPolicies.filter(isAll).map((p) => String(p.subject_id)), + ); + return teams.filter((team) => !covered.has(String(team.id))); + }, [teams, teamPolicies]); + + if (data === null && loading) return ; + if (!data?.success) return ; + + const bucketPill = (policy: QuotaPolicy) => + isAll(policy) ? null : {policy.bucket} traffic; + + return ( +
+

+ Usage is counted per user over each calendar {data.period} (UTC). The + current window resets {fmtDate(data.resets_at)}. A request is refused + once a budget is used up; the request that crosses it still completes. +

+ + {(data.unpriced_models ?? []).length > 0 ? ( +
+

+ Models without a price are invisible to cost limits +

+

+ These were used this {data.period} and recorded at $0:{' '} + {(data.unpriced_models as any[]) + .map((m) => `${m.model_id} (${fmtNumber(m.tokens)} tokens)`) + .join(', ')} + . Use a token limit for them, or declare their rates in the model + catalog. +

+
+ ) : null} + +
+
+

Instance default

+ +
+

+ {instancePolicy + ? `Tokens: ${describeBudget(instancePolicy.token_limit, instancePolicy.token_unlimited, 'tokens')} · Cost: ${describeBudget(instancePolicy.cost_limit_usd, instancePolicy.cost_unlimited, 'cost')}` + : 'No default: users without a team allowance or override are unlimited.'} +

+
+ +
+
+

Team allowances

+ {teamsWithoutPolicy.length > 0 ? ( +
+ + +
+ ) : null} +
+ {teamPolicies.length === 0 ? ( +

+ No team has an allowance. +

+ ) : ( + + + + + Team + Members + Tokens + Cost + Note + Updated + Actions + + + + {teamPolicies.map((policy) => ( + + + + {policy.team_name ?? policy.subject_id} + + {bucketPill(policy)} + + + {fmtNumber(policy.member_count)} + + + + {isAll(policy) ? ( + + ) : null} + + + ))} + +
+ + )} +
+ +
+

User overrides

+ {userPolicies.length === 0 ? ( +

+ No user has an override. Add one from a user's menu on the + Users tab. +

+ ) : ( + + + + + User + Tokens + Cost + Note + Updated + Actions + + + + {userPolicies.map((policy) => ( + + + + {policy.subject_id} + + {bucketPill(policy)} + + + + {isAll(policy) ? ( + + ) : null} + + + ))} + +
+ + )} +
+ + { + if (!open) setEditing(null); + }} + title={editing ? `Quota · ${editing.title}` : 'Quota'} + > + {editing ? ( + { + setEditing(null); + setTeamPick(''); + load(); + }} + /> + ) : null} + +
+ ); +} diff --git a/frontend/src/admin/UserQuotaModal.tsx b/frontend/src/admin/UserQuotaModal.tsx new file mode 100644 index 00000000..f6ce58b2 --- /dev/null +++ b/frontend/src/admin/UserQuotaModal.tsx @@ -0,0 +1,115 @@ +import { useCallback, useEffect, useState } from 'react'; +import { useSelector } from 'react-redux'; + +import adminService from '../api/services/adminService'; +import teamsService from '../api/services/teamsService'; +import { Modal } from '../components/ui/modal'; +import { selectToken } from '../preferences/preferenceSlice'; +import { LoadError, Loading, fmtDate } from './AdminUI'; +import QuotaEditor, { UsageBar } from './QuotaEditor'; +import { + sourceLabel, + type BucketStatus, + type Budget, + type QuotaPolicy, +} from './quotaUtils'; + +/** A user's effective limits and usage, with the editor for their override. */ +export default function UserQuotaModal({ + userId, + onClose, +}: { + userId: string | null; + onClose: () => void; +}) { + const token = useSelector(selectToken); + const [data, setData] = useState(null); + const [teamNames, setTeamNames] = useState>({}); + + const load = useCallback(async () => { + if (!userId) return; + setData(null); + try { + const [res, teamsJson] = await Promise.all([ + adminService.getUserQuota(userId, token), + teamsService.listAll(token).catch(() => ({})), + ]); + setData(await res.json().catch(() => ({ success: false }))); + setTeamNames( + Object.fromEntries( + (teamsJson?.teams ?? []).map((team: any) => [ + String(team.id), + team.name, + ]), + ), + ); + } catch { + setData({ success: false }); + } + }, [userId, token]); + + useEffect(() => { + load(); + }, [load]); + + const overall: BucketStatus | undefined = (data?.effective ?? []).find( + (status: BucketStatus) => status.bucket === 'all', + ); + const override: QuotaPolicy | null = + (data?.policies ?? []).find((p: QuotaPolicy) => p.bucket === 'all') ?? null; + const caption = (budget: Budget) => + sourceLabel( + budget, + budget.source_id ? teamNames[budget.source_id] : undefined, + ); + + return ( + { + if (!open) onClose(); + }} + title={userId ? `Quota · ${userId}` : 'Quota'} + > + {data === null ? ( + + ) : !data.success ? ( + + ) : ( +
+ {overall ? ( +
+ + +

+ Resets {fmtDate(overall.resets_at)} +

+
+ ) : null} +
+

+ User override +

+ +
+
+ )} + + ); +} diff --git a/frontend/src/admin/Users.tsx b/frontend/src/admin/Users.tsx index 13e40d13..d7ae7f8e 100644 --- a/frontend/src/admin/Users.tsx +++ b/frontend/src/admin/Users.tsx @@ -1,5 +1,6 @@ import { Eye, + Gauge, LogOut, ShieldCheck, ShieldOff, @@ -40,6 +41,7 @@ import { fmtNumber, fmtRelative, } from './AdminUI'; +import UserQuotaModal from './UserQuotaModal'; type AdminUser = { user_id: string; @@ -70,6 +72,7 @@ export default function Users() { const [busy, setBusy] = useState(null); const [menuUserId, setMenuUserId] = useState(null); const [detail, setDetail] = useState(null); + const [quotaUserId, setQuotaUserId] = useState(null); const [feedback, setFeedback] = useState<{ ok: boolean; message: string; @@ -157,7 +160,14 @@ export default function Users() { isAdmin: boolean, active: boolean, ): Action[] => { - const acts: Action[] = []; + const acts: Action[] = [ + { + key: 'quota', + label: 'Quota', + icon: Gauge, + perform: () => setQuotaUserId(userId), + }, + ]; if (isAdmin) { acts.push({ key: 'revoke', @@ -433,6 +443,11 @@ export default function Users() { /> ) : null} + setQuotaUserId(null)} + /> + { diff --git a/frontend/src/admin/index.tsx b/frontend/src/admin/index.tsx index 6a147f06..5a21c9f4 100644 --- a/frontend/src/admin/index.tsx +++ b/frontend/src/admin/index.tsx @@ -11,6 +11,7 @@ import { Tabs, TabsList, TabsTrigger } from '../components/ui/tabs'; import Admins from './Admins'; import Audit from './Audit'; import Overview from './Overview'; +import Quotas from './Quotas'; import Usage from './Usage'; import Users from './Users'; @@ -19,6 +20,7 @@ const TABS = [ { key: 'users', label: 'Users', path: '/admin/users' }, { key: 'admins', label: 'Admins', path: '/admin/roles' }, { key: 'usage', label: 'Usage', path: '/admin/usage' }, + { key: 'quotas', label: 'Quotas', path: '/admin/quotas' }, { key: 'audit', label: 'Audit', path: '/admin/audit' }, ]; @@ -63,6 +65,7 @@ export default function Admin() { } /> } /> } /> + } /> } /> } /> diff --git a/frontend/src/admin/quotaUtils.test.ts b/frontend/src/admin/quotaUtils.test.ts new file mode 100644 index 00000000..60e4f468 --- /dev/null +++ b/frontend/src/admin/quotaUtils.test.ts @@ -0,0 +1,120 @@ +import { describe, expect, it } from 'vitest'; + +import { + describeBudget, + formToPolicy, + isEmptyForm, + policyToForm, + sourceLabel, + usagePercent, + type QuotaPolicy, +} from './quotaUtils'; + +const policy = (fields: Partial): QuotaPolicy => ({ + scope: 'user', + subject_id: 'u1', + bucket: 'all', + token_limit: null, + token_unlimited: false, + cost_limit_usd: null, + cost_unlimited: false, + enabled: true, + ...fields, +}); + +describe('policyToForm', () => { + it('starts a missing policy as inherit', () => { + const form = policyToForm(null); + expect(form.tokenMode).toBe('inherit'); + expect(form.costMode).toBe('inherit'); + expect(isEmptyForm(form)).toBe(true); + }); + + it('keeps zero as a limit, not as inherit', () => { + const form = policyToForm(policy({ token_limit: 0, cost_unlimited: true })); + expect(form.tokenMode).toBe('limit'); + expect(form.tokenLimit).toBe('0'); + expect(form.costMode).toBe('unlimited'); + }); +}); + +describe('formToPolicy', () => { + const base = policyToForm(null); + + it('round-trips limits and trims the note', () => { + const result = formToPolicy({ + ...base, + tokenMode: 'limit', + tokenLimit: ' 5000 ', + costMode: 'limit', + costLimit: '2.5', + note: ' trial ', + }); + expect(result).toEqual({ + ok: true, + policy: { + bucket: 'all', + token_limit: 5000, + token_unlimited: false, + cost_limit_usd: 2.5, + cost_unlimited: false, + note: 'trial', + }, + }); + }); + + it('sends unlimited without a limit', () => { + const result = formToPolicy({ + ...base, + tokenMode: 'unlimited', + tokenLimit: '99', + }); + expect(result.ok && result.policy.token_limit).toBeNull(); + expect(result.ok && result.policy.token_unlimited).toBe(true); + }); + + it.each(['', '1.5', '-1', 'abc', '1e3'])( + 'rejects token limit %j', + (tokenLimit) => { + expect(formToPolicy({ ...base, tokenMode: 'limit', tokenLimit }).ok).toBe( + false, + ); + }, + ); + + it.each(['', '-0.01', 'abc', 'Infinity'])( + 'rejects cost limit %j', + (costLimit) => { + expect(formToPolicy({ ...base, costMode: 'limit', costLimit }).ok).toBe( + false, + ); + }, + ); +}); + +describe('usagePercent', () => { + it('handles unlimited, zero and overshoot', () => { + expect(usagePercent(50, null)).toBe(0); + expect(usagePercent(0, 0)).toBe(100); + expect(usagePercent(25, 100)).toBe(25); + expect(usagePercent(500, 100)).toBe(100); + }); +}); + +describe('labels', () => { + it('describes budgets', () => { + expect(describeBudget(null, true, 'tokens')).toBe('Unlimited'); + expect(describeBudget(null, false, 'cost')).toBe('—'); + expect(describeBudget(1000, false, 'tokens')).toContain('tokens'); + }); + + it('names the layer a limit came from', () => { + expect(sourceLabel({ limit: null, used: 0 })).toBe('No limit set'); + expect(sourceLabel({ limit: 1, used: 0, source: 'team' }, 'Eng')).toBe( + 'Team: Eng', + ); + expect(sourceLabel({ limit: 1, used: 0, source: 'default' })).toBe( + 'Plan default', + ); + }); +}); diff --git a/frontend/src/admin/quotaUtils.ts b/frontend/src/admin/quotaUtils.ts new file mode 100644 index 00000000..10038574 --- /dev/null +++ b/frontend/src/admin/quotaUtils.ts @@ -0,0 +1,130 @@ +// Pure helpers behind the quota editor and usage bars. + +export type BudgetMode = 'inherit' | 'limit' | 'unlimited'; + +export type QuotaPolicy = { + scope: 'instance' | 'team' | 'user'; + subject_id: string | null; + bucket: string; + token_limit: number | null; + token_unlimited: boolean; + cost_limit_usd: number | null; + cost_unlimited: boolean; + enabled: boolean; + note?: string | null; + updated_by?: string | null; + updated_at?: string | null; +}; + +export type Budget = { + limit: number | null; + used: number; + source?: string | null; + source_id?: string | null; +}; + +export type BucketStatus = { + bucket: string; + tokens: Budget; + cost: Budget; + resets_at: string; +}; + +export type QuotaForm = { + tokenMode: BudgetMode; + tokenLimit: string; + costMode: BudgetMode; + costLimit: string; + note: string; +}; + +const mode = (limit: number | null, unlimited: boolean): BudgetMode => { + if (unlimited) return 'unlimited'; + return limit === null || limit === undefined ? 'inherit' : 'limit'; +}; + +export function policyToForm(policy?: QuotaPolicy | null): QuotaForm { + return { + tokenMode: policy + ? mode(policy.token_limit, policy.token_unlimited) + : 'inherit', + tokenLimit: policy?.token_limit != null ? String(policy.token_limit) : '', + costMode: policy + ? mode(policy.cost_limit_usd, policy.cost_unlimited) + : 'inherit', + costLimit: + policy?.cost_limit_usd != null ? String(policy.cost_limit_usd) : '', + note: policy?.note ?? '', + }; +} + +export type FormResult = + { ok: true; policy: Record } | { ok: false; error: string }; + +// An empty form (both budgets inherited) is not a policy: the caller deletes instead. +export function isEmptyForm(form: QuotaForm): boolean { + return form.tokenMode === 'inherit' && form.costMode === 'inherit'; +} + +export function formToPolicy(form: QuotaForm): FormResult { + const policy: Record = { + bucket: 'all', + token_limit: null, + token_unlimited: form.tokenMode === 'unlimited', + cost_limit_usd: null, + cost_unlimited: form.costMode === 'unlimited', + note: form.note.trim() || null, + }; + if (form.tokenMode === 'limit') { + const raw = form.tokenLimit.trim(); + if (!/^\d+$/.test(raw)) + return { ok: false, error: 'Token limit must be a whole number.' }; + const tokens = Number(raw); + if (!Number.isSafeInteger(tokens)) + return { ok: false, error: 'Token limit is too large.' }; + policy.token_limit = tokens; + } + if (form.costMode === 'limit') { + const raw = form.costLimit.trim(); + const cost = Number(raw); + if (raw === '' || !Number.isFinite(cost) || cost < 0) + return { ok: false, error: 'Cost limit must be a number, 0 or more.' }; + policy.cost_limit_usd = cost; + } + return { ok: true, policy }; +} + +export function usagePercent(used: number, limit: number | null): number { + if (limit === null || limit === undefined) return 0; + if (limit <= 0) return 100; + return Math.min(100, Math.max(0, (used / limit) * 100)); +} + +export function fmtUsd(value?: number | null): string { + return new Intl.NumberFormat(undefined, { + style: 'currency', + currency: 'USD', + maximumFractionDigits: value != null && value < 1 ? 4 : 2, + }).format(value ?? 0); +} + +export function describeBudget( + limit: number | null, + unlimited: boolean, + kind: 'tokens' | 'cost', +): string { + if (unlimited) return 'Unlimited'; + if (limit === null || limit === undefined) return '—'; + return kind === 'cost' + ? fmtUsd(limit) + : `${new Intl.NumberFormat().format(limit)} tokens`; +} + +export function sourceLabel(budget: Budget, teamName?: string): string { + if (!budget.source) return 'No limit set'; + if (budget.source === 'user') return 'User override'; + if (budget.source === 'team') + return teamName ? `Team: ${teamName}` : 'Team allowance'; + if (budget.source === 'instance') return 'Instance default'; + return 'Plan default'; +} diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 6f5f36e6..89187528 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -2,6 +2,7 @@ const endpoints = { USER: { CONFIG: '/api/config', ME: '/api/user/me', + QUOTA: '/api/user/quota', NEW_TOKEN: '/api/generate_token', OIDC_LOGIN: '/api/auth/oidc/login', OIDC_TOKEN: '/api/auth/oidc/token', @@ -173,6 +174,12 @@ const endpoints = { USAGE: '/api/admin/usage', AUDIT: '/api/admin/audit', DEVICE_AUDIT: '/api/admin/devices/audit', + QUOTAS: '/api/admin/quotas', + QUOTA_INSTANCE: '/api/admin/quotas/instance', + QUOTA_TEAM: (id: string) => + `/api/admin/quotas/teams/${encodeURIComponent(id)}`, + QUOTA_USER: (id: string) => + `/api/admin/quotas/users/${encodeURIComponent(id)}`, }, CONVERSATION: { ANSWER: '/api/answer', diff --git a/frontend/src/api/services/adminService.ts b/frontend/src/api/services/adminService.ts index 7359aed4..adf67e71 100644 --- a/frontend/src/api/services/adminService.ts +++ b/frontend/src/api/services/adminService.ts @@ -10,6 +10,14 @@ const qs = (params: Record): string => { return str ? `?${str}` : ''; }; +export type QuotaScope = 'instance' | 'team' | 'user'; + +const quotaUrl = (scope: QuotaScope, subjectId?: string | null): string => { + if (scope === 'team') return endpoints.ADMIN.QUOTA_TEAM(subjectId ?? ''); + if (scope === 'user') return endpoints.ADMIN.QUOTA_USER(subjectId ?? ''); + return endpoints.ADMIN.QUOTA_INSTANCE; +}; + const adminService = { getOverview: (token: string | null): Promise => apiClient.get(endpoints.ADMIN.OVERVIEW, token), @@ -54,6 +62,23 @@ const adminService = { token: string | null, ): Promise => apiClient.get(`${endpoints.ADMIN.DEVICE_AUDIT}${qs(params)}`, token), + getQuotas: (token: string | null): Promise => + apiClient.get(endpoints.ADMIN.QUOTAS, token), + getUserQuota: (userId: string, token: string | null): Promise => + apiClient.get(endpoints.ADMIN.QUOTA_USER(userId), token), + setQuota: ( + scope: QuotaScope, + subjectId: string | null, + policy: Record, + token: string | null, + ): Promise => apiClient.put(quotaUrl(scope, subjectId), policy, token), + deleteQuota: ( + scope: QuotaScope, + subjectId: string | null, + bucket: string, + token: string | null, + ): Promise => + apiClient.delete(`${quotaUrl(scope, subjectId)}${qs({ bucket })}`, token), }; export default adminService; diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index b9de3bc0..fd9d81f6 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -7,6 +7,8 @@ const userService = { throttledApiClient.get(endpoints.USER.CONFIG, null), getMe: (token: string | null): Promise => apiClient.get(endpoints.USER.ME, token), + getQuota: (token: string | null): Promise => + apiClient.get(endpoints.USER.QUOTA, token), getNewToken: (): Promise => throttledApiClient.get(endpoints.USER.NEW_TOKEN, null), // Token deliberately null: a stale Authorization header must not be able diff --git a/frontend/src/conversation/conversationHandlers.ts b/frontend/src/conversation/conversationHandlers.ts index 4578c8d2..dd519628 100644 --- a/frontend/src/conversation/conversationHandlers.ts +++ b/frontend/src/conversation/conversationHandlers.ts @@ -1,7 +1,10 @@ +import i18n from 'i18next'; + import { baseURL } from '../api/client'; import conversationService from '../api/services/conversationService'; import { Doc } from '../models/misc'; import { Answer, FEEDBACK, RetrievalPayload } from './conversationModels'; +import { isQuotaError, quotaErrorMessage } from './quotaError'; import { ToolCallsType } from './types'; /** @@ -48,7 +51,9 @@ async function _handlePreStreamHttpError( if (text) { try { const parsed = JSON.parse(text); - if (parsed && typeof parsed === 'object') { + if (isQuotaError(parsed)) { + message = quotaErrorMessage(parsed, i18n.t.bind(i18n), i18n.language); + } else if (parsed && typeof parsed === 'object') { message = (typeof parsed.message === 'string' && parsed.message) || (typeof parsed.error === 'string' && parsed.error) || diff --git a/frontend/src/conversation/quotaError.test.ts b/frontend/src/conversation/quotaError.test.ts new file mode 100644 index 00000000..78220511 --- /dev/null +++ b/frontend/src/conversation/quotaError.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from 'vitest'; + +import { isQuotaError, quotaErrorMessage } from './quotaError'; + +const t = ((key: string, values: Record) => + `${key}|${values.used}|${values.limit}|${values.resetsAt}`) as any; + +describe('isQuotaError', () => { + it('matches only the quota error code', () => { + expect(isQuotaError({ error_code: 'quota-exceeded' })).toBe(true); + expect(isQuotaError({ message: 'Exceeding usage limit' })).toBe(false); + expect(isQuotaError(null)).toBe(false); + expect(isQuotaError('quota-exceeded')).toBe(false); + }); +}); + +describe('quotaErrorMessage', () => { + it('formats token budgets as numbers', () => { + const message = quotaErrorMessage( + { + dimension: 'tokens', + usage: 1200000, + limit: 1000000, + resets_at: '2026-10-01T00:00:00+00:00', + }, + t, + 'en-US', + ); + const [key, used, limit, resetsAt] = message.split('|'); + expect(key).toBe('conversation.quotaExceeded.tokens'); + expect([used, limit]).toEqual(['1,200,000', '1,000,000']); + expect(resetsAt).not.toBe(''); + }); + + it('formats cost budgets as dollars', () => { + const message = quotaErrorMessage( + { dimension: 'cost', usage: 5.25, limit: 5 }, + t, + 'en-US', + ); + expect(message).toBe('conversation.quotaExceeded.cost|$5.25|$5.00|'); + }); + + it('tolerates a malformed reset time', () => { + const message = quotaErrorMessage( + { dimension: 'tokens', usage: 1, limit: 1, resets_at: 'soon' }, + t, + 'en-US', + ); + expect(message.endsWith('|')).toBe(true); + }); +}); diff --git a/frontend/src/conversation/quotaError.ts b/frontend/src/conversation/quotaError.ts new file mode 100644 index 00000000..86bf46dd --- /dev/null +++ b/frontend/src/conversation/quotaError.ts @@ -0,0 +1,47 @@ +import type { TFunction } from 'i18next'; + +export type QuotaErrorBody = { + error_code?: string; + dimension?: string; + usage?: number; + limit?: number; + resets_at?: string; +}; + +export function isQuotaError(body: unknown): body is QuotaErrorBody { + return ( + !!body && + typeof body === 'object' && + (body as QuotaErrorBody).error_code === 'quota-exceeded' + ); +} + +/** The chat message for a 429 ``quota-exceeded`` body, in the user's language. */ +export function quotaErrorMessage( + body: QuotaErrorBody, + t: TFunction, + locale?: string, +): string { + const isCost = body.dimension === 'cost'; + const amount = (value?: number) => + isCost + ? new Intl.NumberFormat(locale, { + style: 'currency', + currency: 'USD', + }).format(value ?? 0) + : new Intl.NumberFormat(locale).format(value ?? 0); + const reset = body.resets_at ? new Date(body.resets_at) : null; + const resetsAt = + reset && !Number.isNaN(reset.getTime()) + ? new Intl.DateTimeFormat(locale, { + dateStyle: 'medium', + timeStyle: 'short', + }).format(reset) + : ''; + return t( + isCost + ? 'conversation.quotaExceeded.cost' + : 'conversation.quotaExceeded.tokens', + { used: amount(body.usage), limit: amount(body.limit), resetsAt }, + ); +} diff --git a/frontend/src/locale/de.json b/frontend/src/locale/de.json index ae7d2630..9c6879f7 100644 --- a/frontend/src/locale/de.json +++ b/frontend/src/locale/de.json @@ -375,6 +375,13 @@ "toolCalls": "Werkzeugaufrufe", "runSuccess": "Erfolgsquote", "feedback": "Feedback" + }, + "quota": { + "title": "Ihr Nutzungskontingent", + "resets": "Wird zurückgesetzt: {{resetsAt}}", + "tokens": "Tokens", + "cost": "Kosten", + "usedOf": "{{used}} von {{limit}}" } }, "logs": { @@ -1262,6 +1269,10 @@ "running": "Läuft…", "denied": "Vom Benutzer abgelehnt", "failed": "fehlgeschlagen" + }, + "quotaExceeded": { + "tokens": "Sie haben {{used}} von Ihrem Kontingent von {{limit}} Tokens verbraucht. Es wird am {{resetsAt}} zurückgesetzt.", + "cost": "Sie haben {{used}} von Ihrem Nutzungsbudget von {{limit}} verbraucht. Es wird am {{resetsAt}} zurückgesetzt." } }, "agents": { diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index 93674148..ba269327 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -380,6 +380,13 @@ "toolCalls": "Tool Calls", "runSuccess": "Run Success", "feedback": "Feedback" + }, + "quota": { + "title": "Your usage quota", + "resets": "Resets {{resetsAt}}", + "tokens": "Tokens", + "cost": "Cost", + "usedOf": "{{used}} of {{limit}}" } }, "logs": { @@ -1273,6 +1280,10 @@ "running": "Running…", "denied": "Denied by user", "failed": "failed" + }, + "quotaExceeded": { + "tokens": "You've used {{used}} of your {{limit}} token quota. It resets {{resetsAt}}.", + "cost": "You've used {{used}} of your {{limit}} usage budget. It resets {{resetsAt}}." } }, "agents": { diff --git a/frontend/src/locale/es.json b/frontend/src/locale/es.json index 3f19b13b..b844829b 100644 --- a/frontend/src/locale/es.json +++ b/frontend/src/locale/es.json @@ -375,6 +375,13 @@ "toolCalls": "Llamadas a Herramientas", "runSuccess": "Éxito de Ejecución", "feedback": "Retroalimentación" + }, + "quota": { + "title": "Tu cuota de uso", + "resets": "Se restablece el {{resetsAt}}", + "tokens": "Tokens", + "cost": "Coste", + "usedOf": "{{used}} de {{limit}}" } }, "logs": { @@ -1262,6 +1269,10 @@ "running": "Ejecutando…", "denied": "Denegado por el usuario", "failed": "falló" + }, + "quotaExceeded": { + "tokens": "Has usado {{used}} de tu cuota de {{limit}} tokens. Se restablece el {{resetsAt}}.", + "cost": "Has usado {{used}} de tu presupuesto de uso de {{limit}}. Se restablece el {{resetsAt}}." } }, "agents": { diff --git a/frontend/src/locale/jp.json b/frontend/src/locale/jp.json index f60227d1..7c5aec0b 100644 --- a/frontend/src/locale/jp.json +++ b/frontend/src/locale/jp.json @@ -375,6 +375,13 @@ "toolCalls": "ツール呼び出し", "runSuccess": "実行成功率", "feedback": "フィードバック" + }, + "quota": { + "title": "利用クォータ", + "resets": "{{resetsAt}} にリセット", + "tokens": "トークン", + "cost": "コスト", + "usedOf": "{{used}} / {{limit}}" } }, "logs": { @@ -1262,6 +1269,10 @@ "running": "実行中…", "denied": "ユーザーによって拒否されました", "failed": "失敗" + }, + "quotaExceeded": { + "tokens": "トークンクォータ {{limit}} のうち {{used}} を使用しました。{{resetsAt}} にリセットされます。", + "cost": "利用予算 {{limit}} のうち {{used}} を使用しました。{{resetsAt}} にリセットされます。" } }, "agents": { diff --git a/frontend/src/locale/ru.json b/frontend/src/locale/ru.json index 878451ba..69c0c25f 100644 --- a/frontend/src/locale/ru.json +++ b/frontend/src/locale/ru.json @@ -375,6 +375,13 @@ "toolCalls": "Вызовы инструментов", "runSuccess": "Успешность запусков", "feedback": "Обратная связь" + }, + "quota": { + "title": "Ваша квота использования", + "resets": "Сброс: {{resetsAt}}", + "tokens": "Токены", + "cost": "Стоимость", + "usedOf": "{{used}} из {{limit}}" } }, "logs": { @@ -1282,6 +1289,10 @@ "running": "Выполняется…", "denied": "Отклонено пользователем", "failed": "не удалось" + }, + "quotaExceeded": { + "tokens": "Вы использовали {{used}} из квоты в {{limit}} токенов. Квота сбросится {{resetsAt}}.", + "cost": "Вы использовали {{used}} из бюджета в {{limit}}. Бюджет сбросится {{resetsAt}}." } }, "agents": { diff --git a/frontend/src/locale/zh-TW.json b/frontend/src/locale/zh-TW.json index 5323803d..4de54e9d 100644 --- a/frontend/src/locale/zh-TW.json +++ b/frontend/src/locale/zh-TW.json @@ -375,6 +375,13 @@ "toolCalls": "工具呼叫", "runSuccess": "執行成功率", "feedback": "回饋" + }, + "quota": { + "title": "您的用量配額", + "resets": "{{resetsAt}} 重設", + "tokens": "權杖", + "cost": "費用", + "usedOf": "{{used}} / {{limit}}" } }, "logs": { @@ -1262,6 +1269,10 @@ "running": "執行中…", "denied": "已被使用者拒絕", "failed": "失敗" + }, + "quotaExceeded": { + "tokens": "您已使用 {{limit}} 權杖配額中的 {{used}}。配額將於 {{resetsAt}} 重設。", + "cost": "您已使用 {{limit}} 用量預算中的 {{used}}。預算將於 {{resetsAt}} 重設。" } }, "agents": { diff --git a/frontend/src/locale/zh.json b/frontend/src/locale/zh.json index 93ba0088..6a542506 100644 --- a/frontend/src/locale/zh.json +++ b/frontend/src/locale/zh.json @@ -375,6 +375,13 @@ "toolCalls": "工具调用", "runSuccess": "运行成功率", "feedback": "反馈" + }, + "quota": { + "title": "您的用量配额", + "resets": "{{resetsAt}} 重置", + "tokens": "令牌", + "cost": "费用", + "usedOf": "{{used}} / {{limit}}" } }, "logs": { @@ -1262,6 +1269,10 @@ "running": "正在运行…", "denied": "已被用户拒绝", "failed": "失败" + }, + "quotaExceeded": { + "tokens": "您已使用 {{limit}} 令牌配额中的 {{used}}。配额将于 {{resetsAt}} 重置。", + "cost": "您已使用 {{limit}} 用量预算中的 {{used}}。预算将于 {{resetsAt}} 重置。" } }, "agents": { diff --git a/frontend/src/settings/Analytics.tsx b/frontend/src/settings/Analytics.tsx index 1183f2af..4eaccd5c 100644 --- a/frontend/src/settings/Analytics.tsx +++ b/frontend/src/settings/Analytics.tsx @@ -26,6 +26,7 @@ import { useDarkTheme, useLoaderState } from '../hooks'; import { selectToken } from '../preferences/preferenceSlice'; import { htmlLegendPlugin } from '../utils/chartUtils'; import { formatDate } from '../utils/dateTimeUtils'; +import UsageQuota from './components/UsageQuota'; /** * Resolve a CSS custom property on `:root` to a concrete color string. @@ -377,6 +378,7 @@ export default function Analytics({ agentId }: AnalyticsProps) { return (
+ {agentId ? null : }

{t('settings.analytics.subtitle')} @@ -412,7 +414,7 @@ export default function Analytics({ agentId }: AnalyticsProps) {

{card.label}

diff --git a/frontend/src/settings/components/UsageQuota.tsx b/frontend/src/settings/components/UsageQuota.tsx new file mode 100644 index 00000000..7cb9ceff --- /dev/null +++ b/frontend/src/settings/components/UsageQuota.tsx @@ -0,0 +1,130 @@ +import { useEffect, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useSelector } from 'react-redux'; + +import userService from '../../api/services/userService'; +import { selectToken } from '../../preferences/preferenceSlice'; + +type Budget = { limit: number | null; used: number }; +type Bucket = { + bucket: string; + tokens: Budget; + cost: Budget; + resets_at: string; +}; + +function Meter({ + label, + budget, + format, +}: { + label: string; + budget: Budget; + format: (value: number) => string; +}) { + const { t } = useTranslation(); + if (budget.limit === null) return null; + const percent = + budget.limit <= 0 + ? 100 + : Math.min(100, Math.max(0, (budget.used / budget.limit) * 100)); + const tone = + percent >= 100 + ? 'bg-red-500' + : percent >= 80 + ? 'bg-amber-500' + : 'bg-[#7D54D1]'; + return ( +

+
+ {label} + + {t('settings.analytics.quota.usedOf', { + used: format(budget.used), + limit: format(budget.limit), + })} + +
+
+
+
+
+ ); +} + +/** The caller's usage against the quota an admin set; renders nothing when unlimited. */ +export default function UsageQuota() { + const { t, i18n } = useTranslation(); + const token = useSelector(selectToken); + const [bucket, setBucket] = useState(null); + + useEffect(() => { + let cancelled = false; + userService + .getQuota(token) + .then((res: Response) => (res.ok ? res.json() : null)) + .then((json: { buckets?: Bucket[] } | null) => { + if (cancelled) return; + const buckets = json?.buckets ?? []; + setBucket( + buckets.find((b) => b.bucket === 'all') ?? buckets[0] ?? null, + ); + }) + .catch(() => undefined); + return () => { + cancelled = true; + }; + }, [token]); + + if (!bucket) return null; + + const number = new Intl.NumberFormat(i18n.language); + const usd = new Intl.NumberFormat(i18n.language, { + style: 'currency', + currency: 'USD', + }); + const reset = new Date(bucket.resets_at); + const resetsAt = Number.isNaN(reset.getTime()) + ? '' + : new Intl.DateTimeFormat(i18n.language, { + dateStyle: 'medium', + timeStyle: 'short', + }).format(reset); + + return ( +
+
+

+ {t('settings.analytics.quota.title')} +

+ {resetsAt ? ( +

+ {t('settings.analytics.quota.resets', { resetsAt })} +

+ ) : null} +
+
+ number.format(value)} + /> + usd.format(value)} + /> +
+
+ ); +} From 1eacfdd3d00fa514ff39aa30fe4c2b5582229e8a Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 12:03:02 +0100 Subject: [PATCH 08/14] docs: usage quotas How the instance default, team allowances and user overrides resolve (including users in several teams), the quota window, who is charged for agent traffic, how cost budgets price models and what happens to unpriced ones, and the admin and user API. --- .env-template | 6 ++ docs/content/Deploying/Access-Control.mdx | 4 +- docs/content/Deploying/Usage-Quotas.mdx | 109 ++++++++++++++++++++++ docs/content/Deploying/_meta.js | 4 + 4 files changed, 122 insertions(+), 1 deletion(-) create mode 100644 docs/content/Deploying/Usage-Quotas.mdx diff --git a/.env-template b/.env-template index 5e868f5e..7f3c05a6 100644 --- a/.env-template +++ b/.env-template @@ -109,3 +109,9 @@ MICROSOFT_AUTHORITY=https://{tenantId}.ciamlogin.com/{tenantId} # PAT_MAX_LIFETIME_DAYS=365 # PAT_ALLOW_NON_EXPIRING=false # PAT_MAX_PER_USER=25 + +# Usage quotas (set limits in Admin → Quotas). Usage is counted per calendar +# day, week or month in UTC. Models without a declared price are recorded at $0 +# unless a fallback [input, output] USD rate per 1M tokens is given. +# QUOTA_PERIOD=month +# QUOTA_UNPRICED_RATE_PER_MILLION=[0.5, 1.5] diff --git a/docs/content/Deploying/Access-Control.mdx b/docs/content/Deploying/Access-Control.mdx index 1c452004..dc238349 100644 --- a/docs/content/Deploying/Access-Control.mdx +++ b/docs/content/Deploying/Access-Control.mdx @@ -88,6 +88,7 @@ Admins get a dashboard backed by a REST surface under `/api/admin` (every endpoi | `GET` | `/api/admin/audit` | Authentication/admin audit feed. | | `GET` | `/api/admin/devices/audit` | Remote-device audit feed. | | `GET` | `/api/admin/teams` | Instance-wide oversight of all teams. | +| `GET` `PUT` `DELETE` | `/api/admin/quotas/...` | [Usage quotas](/Deploying/Usage-Quotas) for the instance, teams and users. | Deactivating a user via the dashboard works for any auth type, while OIDC deployments can also offboard through [SCIM](/Deploying/OIDC-SSO#scim-user-provisioning). Both revoke live sessions immediately. @@ -141,9 +142,10 @@ Sharing rules: ## Audit log -Access-control actions are appended to the `auth_events` table alongside the [authentication events](/Deploying/OIDC-SSO#login-auditing). This includes admin actions — `admin_user_activated` / `admin_user_deactivated`, `admin_sessions_revoked`, `role_granted` / `role_revoked` (with `metadata.source` = `manual` or `oidc_group`) — and team events (`team.create`, `team.member_add`, `team.member_role`, `team.member_remove`, `team.share`, `team.unshare`, `team.transfer_owner`, `team.delete`). The acting admin is recorded in the event metadata. +Access-control actions are appended to the `auth_events` table alongside the [authentication events](/Deploying/OIDC-SSO#login-auditing). This includes admin actions — `admin_user_activated` / `admin_user_deactivated`, `admin_sessions_revoked`, `role_granted` / `role_revoked` (with `metadata.source` = `manual` or `oidc_group`), `quota_policy_set` / `quota_policy_deleted` — and team events (`team.create`, `team.member_add`, `team.member_role`, `team.member_remove`, `team.share`, `team.unshare`, `team.transfer_owner`, `team.delete`). The acting admin is recorded in the event metadata. ## Related - [SSO with OIDC](/Deploying/OIDC-SSO) — sign-in, group allowlists, and the `auth_events` table. +- [Usage Quotas](/Deploying/Usage-Quotas) — token and cost limits per user and per team. - [App Configuration](/Deploying/DocsGPT-Settings) — the full settings reference. diff --git a/docs/content/Deploying/Usage-Quotas.mdx b/docs/content/Deploying/Usage-Quotas.mdx new file mode 100644 index 00000000..701b9c83 --- /dev/null +++ b/docs/content/Deploying/Usage-Quotas.mdx @@ -0,0 +1,109 @@ +--- +title: Usage Quotas +description: Cap how many tokens or dollars each user may spend per day, week or month, with an instance default, per-team allowances and per-user overrides. +--- + +import { Callout } from 'nextra/components' + +# Usage Quotas + +An instance admin can limit how much each user spends on language models. A quota has two independent budgets: + +- **Tokens** — prompt plus generated tokens. Works for every model, including local ones. +- **Cost (USD)** — tokens priced at the model's catalog rate. Only sees models that declare a price. + +Set either, both or neither. Quotas are managed from **Admin → Quotas**, or through the [API](#api). With no quota set, nothing is limited. + +## Layers + +Limits are set at three layers. For each budget, the first layer that says something wins: + +1. **User override** — one user's own limit. +2. **Team allowance** — what each member of a team gets. +3. **Instance default** — everyone else. + +At each layer a budget is either *not set* (defer to the next layer), a *limit*, or *unlimited*. A limit of `0` blocks the user. The two budgets resolve separately, so a user's token limit can come from their team while their cost limit comes from the instance default. + +### Teams + +A team allowance is **per member**, not a pool the team shares: if the allowance is 2M tokens, each member may use 2M. + +A user in several teams gets the **most generous** allowance among them, and allowances are never added together. Usage is always counted per user, whichever teams they belong to. To hold one person below their team's allowance, give them a user override. + + +Team membership can change without an instance admin — team admins, OIDC group sync and SCIM all add members — so joining a team can only raise a user's allowance to what you granted that team, never lower it. Only instance admins set allowances; team admins cannot. + + +## Windows and enforcement + +Usage is counted over a calendar window in UTC, chosen for the whole instance with [`QUOTA_PERIOD`](/Deploying/Settings-Reference#quotas): `day` (from 00:00), `week` (from Monday) or `month` (from the 1st, the default). Windows are worked out when a request arrives, so there is no reset job to run. + +The quota is checked **before** a request starts. The request that crosses a limit completes; the next one is refused with HTTP `429`: + +```json +{ + "success": false, + "error_code": "quota-exceeded", + "message": "Usage quota reached (1,000,000 of 1,000,000 tokens). It resets at 2026-10-01T00:00:00+00:00.", + "dimension": "tokens", + "unit": "tokens", + "usage": 1000000, + "limit": 1000000, + "bucket": "all", + "source": "instance", + "resets_at": "2026-10-01T00:00:00+00:00" +} +``` + +The response carries a `Retry-After` header. The check covers chat, the agent and OpenAI-compatible APIs, scheduled runs (recorded as `budget_exceeded`) and webhook runs. If the quota check itself fails, the request is allowed. + +Who is charged: + +| Traffic | Charged to | +| --- | --- | +| Chat without an agent | The user | +| A user's own agent, its API key, webhooks and schedules | The agent's owner | +| An agent shared with the user | The user | + +Per-agent token and request limits still apply on top of the owner's quota. + +Users with a quota see their usage and the reset time under **Settings → Analytics**. + +## Pricing + +Cost budgets use the rates in the [model catalog](/Models/cloud-providers), in USD per million tokens: + +```yaml +models: + - id: my-model + input_cost_per_million: 3.0 + output_cost_per_million: 15.0 + cached_input_cost_per_million: 0.3 # optional, prompt-cache reads + cache_write_cost_per_million: 3.75 # optional, prompt-cache writes +``` + +The built-in catalogs ship list prices for hosted models. Override or add rates by dropping a YAML with the same model `id` into `MODELS_CONFIG_DIR`. The cost of each call is stored with its usage row when the call is made, so later price changes do not rewrite history. + + +A model with no declared price is recorded at $0, so a cost budget cannot see it. The Quotas tab lists such models once they have been used. Either limit them with a token budget, declare their rates, or set [`QUOTA_UNPRICED_RATE_PER_MILLION`](/Deploying/Settings-Reference#quotas) to charge a fallback rate. Models a user adds with their own API key are always $0, but their tokens still count. + + +## API + +Every admin endpoint requires the admin role, and every change is written to the [audit log](/Deploying/Access-Control#audit-log) as `quota_policy_set` or `quota_policy_deleted`. + +| Method | Path | Description | +| --- | --- | --- | +| `GET` | `/api/admin/quotas` | All policies by layer, the current window, and used models without a price. | +| `PUT` `DELETE` | `/api/admin/quotas/instance` | The instance default. | +| `GET` `PUT` `DELETE` | `/api/admin/quotas/teams/` | A team's per-member allowance. | +| `GET` `PUT` `DELETE` | `/api/admin/quotas/users/` | A user's override. `GET` also returns the limits the user ends up with, the layer each came from, and their usage. | +| `GET` | `/api/user/quota` | The caller's own limits, usage and reset time. | + +A `PUT` body sets, per budget, a limit or the unlimited flag; leave both out to defer to the next layer: + +```json +{ "token_limit": 2000000, "cost_unlimited": true, "note": "Research team" } +``` + +`bucket` (default `all`) narrows a policy to `direct` traffic (chat without an agent) or `agent` traffic (anything that runs through an agent). A request must fit both its own bucket and `all`. The dashboard edits `all`. diff --git a/docs/content/Deploying/_meta.js b/docs/content/Deploying/_meta.js index 105b9efd..9afb460d 100644 --- a/docs/content/Deploying/_meta.js +++ b/docs/content/Deploying/_meta.js @@ -15,6 +15,10 @@ export default { "title": "👥 Access Control & Teams", "href": "/Deploying/Access-Control" }, + "Usage-Quotas": { + "title": "📊 Usage Quotas", + "href": "/Deploying/Usage-Quotas" + }, "Docker-Deploying": { "title": "🛳️ Docker Setup", "href": "/Deploying/Docker-Deploying" From 6db9014dfb146b37f65dcaafcedba0ce74ed8ce5 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 12:14:14 +0100 Subject: [PATCH 09/14] fix(quotas): list unpriced models by recorded cost; integer token limits in status The unpriced-model notice asked the live registry whether a model has a price, so a priced model whose provider was later disabled showed up as unpriced. It now lists models whose calls this period were all recorded at $0. Token limits are serialized as integers. --- docsgpt/api/admin/quotas.py | 8 +++++--- docsgpt/quotas/service.py | 13 +++++++------ tests/api/test_quota_endpoints.py | 13 ++++++++----- tests/quotas/test_service.py | 2 +- 4 files changed, 21 insertions(+), 15 deletions(-) diff --git a/docsgpt/api/admin/quotas.py b/docsgpt/api/admin/quotas.py index 8210ba8d..58b49387 100644 --- a/docsgpt/api/admin/quotas.py +++ b/docsgpt/api/admin/quotas.py @@ -151,13 +151,15 @@ def _delete_policy(scope: str, subject_id: Optional[str]): def _unpriced_models(conn) -> list[dict]: - """Catalog models used this period that no cost limit can see.""" + """Models used this period whose calls were all recorded at $0 for want of a price.""" start, _ = window_bounds(settings.QUOTA_PERIOD) return [ row for row in TokenUsageRepository(conn).tokens_by_model(start=start) - # BYOM ids are UUIDs; those calls are $0 by design, not by omission. - if not looks_like_uuid(row["model_id"]) and not is_priced(row["model_id"]) + # Judged by what was recorded, so a priced model whose provider has since + # been disabled is not listed. BYOM ids are UUIDs and $0 by design; a + # model explicitly priced at $0 is free, not unpriced. + if row["cost"] == 0 and not looks_like_uuid(row["model_id"]) and not is_priced(row["model_id"]) ] diff --git a/docsgpt/quotas/service.py b/docsgpt/quotas/service.py index 7659fe1d..fff7178a 100644 --- a/docsgpt/quotas/service.py +++ b/docsgpt/quotas/service.py @@ -43,18 +43,19 @@ class BucketStatus: def to_dict(self) -> dict: """Return the JSON shape shared by the admin and user quota endpoints.""" - def budget(limit: ResolvedLimit, used: float) -> dict: + def budget(limit: Optional[float], resolved: ResolvedLimit, used: float) -> dict: return { - "limit": limit.limit, + "limit": limit, "used": used, - "source": limit.source, - "source_id": limit.source_id, + "source": resolved.source, + "source_id": resolved.source_id, } + tokens, cost = self.limits.tokens, self.limits.cost return { "bucket": self.bucket, - "tokens": budget(self.limits.tokens, self.tokens_used), - "cost": budget(self.limits.cost, round(self.cost_used, 6)), + "tokens": budget(None if tokens.unlimited else int(tokens.limit), tokens, self.tokens_used), + "cost": budget(cost.limit, cost, round(self.cost_used, 6)), "resets_at": self.resets_at.isoformat(), } diff --git a/tests/api/test_quota_endpoints.py b/tests/api/test_quota_endpoints.py index 44c424db..c7337d1b 100644 --- a/tests/api/test_quota_endpoints.py +++ b/tests/api/test_quota_endpoints.py @@ -222,7 +222,8 @@ class TestUserPolicy: body = _body(client.get("/api/admin/quotas/users/u1")) overall = body["effective"][0] assert overall["bucket"] == "all" - assert overall["tokens"] == {"limit": 900.0, "used": 40, "source": "team", "source_id": big} + assert overall["tokens"] == {"limit": 900, "used": 40, "source": "team", "source_id": big} + assert isinstance(overall["tokens"]["limit"], int) assert overall["cost"] == {"limit": 1.0, "used": 0.25, "source": "instance", "source_id": None} assert body["policies"] == [] @@ -243,12 +244,14 @@ class TestUserPolicy: class TestUnpricedModels: - def test_lists_used_catalog_models_without_a_price(self, client, db): + def test_lists_models_recorded_at_zero_for_want_of_a_price(self, client, db): usage = TokenUsageRepository(db) usage.insert(user_id="u1", prompt_tokens=10, model_id="local-llama") - usage.insert(user_id="u1", prompt_tokens=5, model_id="claude-haiku-4-5", cost=0.1) + # Priced when called; its provider may be disabled by now. + usage.insert(user_id="u1", prompt_tokens=5, model_id="retired-priced-model", cost=0.1) + usage.insert(user_id="u1", prompt_tokens=3, model_id="free-model") usage.insert(user_id="u1", prompt_tokens=7, model_id="7d0c1a52-2f5e-4c53-9a0e-111111111111") - with _admin(), patch("docsgpt.api.admin.quotas.is_priced", lambda m: m == "claude-haiku-4-5"): + with _admin(), patch("docsgpt.api.admin.quotas.is_priced", lambda m: m == "free-model"): unpriced = _body(client.get("/api/admin/quotas"))["unpriced_models"] assert unpriced == [{"model_id": "local-llama", "tokens": 10, "cost": 0.0}] @@ -266,7 +269,7 @@ class TestMyQuota: body = _body(client.get("/api/user/quota")) (bucket,) = body["buckets"] assert bucket["bucket"] == "all" - assert bucket["tokens"] == {"limit": 100.0, "used": 30} + assert bucket["tokens"] == {"limit": 100, "used": 30} assert bucket["cost"] == {"limit": None, "used": 0.0} assert "source" not in json.dumps(body) and "secret" not in json.dumps(body) diff --git a/tests/quotas/test_service.py b/tests/quotas/test_service.py index 0a365799..13975bb3 100644 --- a/tests/quotas/test_service.py +++ b/tests/quotas/test_service.py @@ -210,7 +210,7 @@ class TestStatusAndPayload: (status,) = QuotaService.status("u1", now=NOW) assert status.to_dict() == { "bucket": "all", - "tokens": {"limit": 100.0, "used": 40, "source": "instance", "source_id": None}, + "tokens": {"limit": 100, "used": 40, "source": "instance", "source_id": None}, "cost": {"limit": 2.0, "used": 0.5, "source": "instance", "source_id": None}, "resets_at": "2026-10-01T00:00:00+00:00", } From 69f55b74cb0ae069ac5699002412f557ad726d38 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 12:15:42 +0100 Subject: [PATCH 10/14] refactor(quotas): validate policy bodies without exception text in responses Validation problems are returned as values rather than raised and echoed with str(exc), and a huge integer limit is rejected as out of range instead of overflowing. Tests no longer call mutating endpoints inside asserts. --- docsgpt/api/admin/quotas.py | 111 +++++++++--------- tests/api/test_quota_endpoints.py | 12 +- .../db/repositories/test_quota_policies.py | 7 +- 3 files changed, 65 insertions(+), 65 deletions(-) diff --git a/docsgpt/api/admin/quotas.py b/docsgpt/api/admin/quotas.py index 58b49387..0235f36b 100644 --- a/docsgpt/api/admin/quotas.py +++ b/docsgpt/api/admin/quotas.py @@ -30,10 +30,8 @@ from docsgpt.storage.db.session import db_readonly, db_session _MAX_TOKEN_LIMIT = 2**62 _MAX_COST_LIMIT = 99_999_999.0 _MAX_NOTE_LENGTH = 500 - - -class _BadPolicy(ValueError): - """The request body does not describe a valid policy.""" +_FLAG_DEFAULTS = {"token_unlimited": False, "cost_unlimited": False, "enabled": True} +_BUCKET_MESSAGE = f"bucket must be one of: {', '.join(BUCKETS)}" def _policy_json(row: dict) -> dict: @@ -57,60 +55,59 @@ def _error(message: str, status: int): return make_response(jsonify({"success": False, "message": message}), status) -def _bucket(value: Any) -> str: +def _limit_error(value: Any, name: str, whole: bool, maximum: float) -> Optional[str]: + """Return why ``value`` is not a valid limit, or ``None``.""" if value is None: - return "all" - if value not in BUCKETS: - raise _BadPolicy(f"bucket must be one of: {', '.join(BUCKETS)}") - return value + return None + number = (int,) if whole else (int, float) + if isinstance(value, bool) or not isinstance(value, number): + return f"{name} must be a {'whole number' if whole else 'number'} or null" + # Range first: ``isfinite`` overflows on an int too large for a float. + if not 0 <= value <= maximum or not math.isfinite(value): + return f"{name} is out of range" + return None -def _flag(data: dict, key: str, default: bool) -> bool: - value = data.get(key, default) - if not isinstance(value, bool): - raise _BadPolicy(f"{key} must be a boolean") - return value +def _parse_policy(data: Any) -> tuple[Optional[dict], Optional[str]]: + """Validate a policy body. - -def _parse_policy(data: Any) -> dict: - """Validate a policy body into ``QuotaPoliciesRepository.upsert`` kwargs.""" + Returns: + ``(fields, None)`` with ``QuotaPoliciesRepository.upsert`` kwargs, or + ``(None, message)`` describing the first problem. + """ if not isinstance(data, dict): - raise _BadPolicy("Body must be a JSON object") - token_limit = data.get("token_limit") - if token_limit is not None: - if isinstance(token_limit, bool) or not isinstance(token_limit, int): - raise _BadPolicy("token_limit must be a whole number or null") - if not 0 <= token_limit <= _MAX_TOKEN_LIMIT: - raise _BadPolicy("token_limit is out of range") - cost_limit = data.get("cost_limit_usd") - if cost_limit is not None: - if isinstance(cost_limit, bool) or not isinstance(cost_limit, (int, float)): - raise _BadPolicy("cost_limit_usd must be a number or null") - if not math.isfinite(cost_limit) or not 0 <= cost_limit <= _MAX_COST_LIMIT: - raise _BadPolicy("cost_limit_usd is out of range") - cost_limit = round(float(cost_limit), 4) - token_unlimited = _flag(data, "token_unlimited", False) - cost_unlimited = _flag(data, "cost_unlimited", False) - if token_unlimited and token_limit is not None: - raise _BadPolicy("Set token_limit or token_unlimited, not both") - if cost_unlimited and cost_limit is not None: - raise _BadPolicy("Set cost_limit_usd or cost_unlimited, not both") - if token_limit is None and cost_limit is None and not token_unlimited and not cost_unlimited: - raise _BadPolicy("Set a limit or mark a budget unlimited; delete the policy to remove it") + return None, "Body must be a JSON object" + token_limit, cost_limit = data.get("token_limit"), data.get("cost_limit_usd") + problem = _limit_error(token_limit, "token_limit", True, _MAX_TOKEN_LIMIT) or _limit_error( + cost_limit, "cost_limit_usd", False, _MAX_COST_LIMIT + ) + if problem: + return None, problem + flags = {key: data.get(key, default) for key, default in _FLAG_DEFAULTS.items()} + for key, value in flags.items(): + if not isinstance(value, bool): + return None, f"{key} must be a boolean" + bucket = data.get("bucket", "all") + if bucket not in BUCKETS: + return None, _BUCKET_MESSAGE note = data.get("note") - if note is not None: - if not isinstance(note, str): - raise _BadPolicy("note must be a string") - note = note.strip()[:_MAX_NOTE_LENGTH] or None + if note is not None and not isinstance(note, str): + return None, "note must be a string" + if flags["token_unlimited"] and token_limit is not None: + return None, "Set token_limit or token_unlimited, not both" + if flags["cost_unlimited"] and cost_limit is not None: + return None, "Set cost_limit_usd or cost_unlimited, not both" + if token_limit is None and cost_limit is None and not flags["token_unlimited"] and not flags["cost_unlimited"]: + return None, "Set a limit or mark a budget unlimited; delete the policy to remove it" return { - "bucket": _bucket(data.get("bucket")), + "bucket": bucket, "token_limit": token_limit, - "token_unlimited": token_unlimited, - "cost_limit_usd": cost_limit, - "cost_unlimited": cost_unlimited, - "enabled": _flag(data, "enabled", True), - "note": note, - } + "token_unlimited": flags["token_unlimited"], + "cost_limit_usd": round(float(cost_limit), 4) if cost_limit is not None else None, + "cost_unlimited": flags["cost_unlimited"], + "enabled": flags["enabled"], + "note": (note.strip()[:_MAX_NOTE_LENGTH] or None) if note else None, + }, None def _audit(conn, event: str, scope: str, subject_id: Optional[str], detail: dict) -> None: @@ -126,10 +123,9 @@ def _audit(conn, event: str, scope: str, subject_id: Optional[str], detail: dict def _put_policy(scope: str, subject_id: Optional[str]): - try: - fields = _parse_policy(request.get_json(silent=True)) - except _BadPolicy as exc: - return _error(str(exc), 400) + fields, problem = _parse_policy(request.get_json(silent=True)) + if fields is None: + return _error(problem or "Invalid policy", 400) with db_session() as conn: row = QuotaPoliciesRepository(conn).upsert( scope=scope, subject_id=subject_id, actor=_actor(), **fields @@ -139,10 +135,9 @@ def _put_policy(scope: str, subject_id: Optional[str]): def _delete_policy(scope: str, subject_id: Optional[str]): - try: - bucket = _bucket(request.args.get("bucket")) if "bucket" in request.args else None - except _BadPolicy as exc: - return _error(str(exc), 400) + bucket = request.args.get("bucket") + if bucket is not None and bucket not in BUCKETS: + return _error(_BUCKET_MESSAGE, 400) with db_session() as conn: deleted = QuotaPoliciesRepository(conn).delete(scope, subject_id, bucket) if deleted: diff --git a/tests/api/test_quota_endpoints.py b/tests/api/test_quota_endpoints.py index c7337d1b..69224f9c 100644 --- a/tests/api/test_quota_endpoints.py +++ b/tests/api/test_quota_endpoints.py @@ -125,9 +125,10 @@ class TestInstancePolicy: ] assert overview["period"] == "month" - assert _body(client.delete("/api/admin/quotas/instance?bucket=agent"))["deleted"] == 1 - assert _body(client.delete("/api/admin/quotas/instance"))["deleted"] == 1 - assert _body(client.get("/api/admin/quotas"))["instance"] == [] + one_bucket = _body(client.delete("/api/admin/quotas/instance?bucket=agent")) + the_rest = _body(client.delete("/api/admin/quotas/instance")) + remaining = _body(client.get("/api/admin/quotas"))["instance"] + assert (one_bucket["deleted"], the_rest["deleted"], remaining) == (1, 1, []) def test_writes_are_audited(self, client, db): with _admin(): @@ -158,6 +159,8 @@ class TestInstancePolicy: {"token_limit": True}, {"token_limit": "10"}, {"token_limit": 2**63}, + {"token_limit": 10**400}, + {"cost_limit_usd": 10**400}, {"cost_limit_usd": -0.01}, {"cost_limit_usd": "5"}, {"cost_limit_usd": float("inf")}, @@ -178,7 +181,8 @@ class TestInstancePolicy: def test_unknown_bucket_on_delete(self, client, db): with _admin(): - assert client.delete("/api/admin/quotas/instance?bucket=nope").status_code == 400 + resp = client.delete("/api/admin/quotas/instance?bucket=nope") + assert resp.status_code == 400 def test_zero_is_accepted_as_a_block(self, client, db): with _admin(): diff --git a/tests/storage/db/repositories/test_quota_policies.py b/tests/storage/db/repositories/test_quota_policies.py index 25396210..83a9302c 100644 --- a/tests/storage/db/repositories/test_quota_policies.py +++ b/tests/storage/db/repositories/test_quota_policies.py @@ -137,7 +137,8 @@ class TestDelete: repo.upsert(scope="user", subject_id="u1", token_limit=1) repo.upsert(scope="user", subject_id="u1", bucket="agent", token_limit=2) repo.upsert(scope="user", subject_id="u2", token_limit=3) - assert repo.delete("user", "u1", "agent") == 1 - assert repo.delete("user", "u1", "agent") == 0 - assert repo.delete("user", "u1") == 1 + first = repo.delete("user", "u1", "agent") + again = repo.delete("user", "u1", "agent") + rest = repo.delete("user", "u1") + assert (first, again, rest) == (1, 0, 1) assert repo.get("user", "u2") is not None From 4bd259fc091932a6ee0c017bb3595111089b5ea6 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 12:44:26 +0100 Subject: [PATCH 11/14] fix(quotas): address review: resume claims, agent bucket rule, UI races - A tool continuation refused for usage now releases the resume claim it took; before, retries got a 409 until the stale claim was reverted. - Agent traffic is any row with an agent key or an agent id, so keyless agents and workflow nodes count toward the agent bucket, not direct. - The user quota modal discards responses for a previously opened user. - The usage meter shows every limited bucket, not only 'all'. - Restore the class separator on the analytics stat card that a formatter run removed, and align the OpenRouter DeepSeek description with its rates. --- docs/content/Deploying/Usage-Quotas.mdx | 2 +- docsgpt/agents/headless_runner.py | 4 +- docsgpt/api/answer/routes/answer.py | 4 +- docsgpt/api/answer/routes/base.py | 26 +++++++ docsgpt/api/answer/routes/stream.py | 6 +- docsgpt/api/v1/routes.py | 8 ++- docsgpt/core/models/openrouter.yaml | 2 +- .../storage/db/repositories/token_usage.py | 9 +-- frontend/src/admin/UserQuotaModal.tsx | 15 +++-- frontend/src/locale/de.json | 6 +- frontend/src/locale/en.json | 6 +- frontend/src/locale/es.json | 6 +- frontend/src/locale/jp.json | 6 +- frontend/src/locale/ru.json | 6 +- frontend/src/locale/zh-TW.json | 6 +- frontend/src/locale/zh.json | 6 +- frontend/src/settings/Analytics.tsx | 2 +- .../src/settings/components/UsageQuota.tsx | 50 ++++++++------ tests/quotas/test_enforcement.py | 67 +++++++++++++++++++ .../db/repositories/test_token_usage.py | 10 ++- 20 files changed, 202 insertions(+), 45 deletions(-) diff --git a/docs/content/Deploying/Usage-Quotas.mdx b/docs/content/Deploying/Usage-Quotas.mdx index 701b9c83..c7391625 100644 --- a/docs/content/Deploying/Usage-Quotas.mdx +++ b/docs/content/Deploying/Usage-Quotas.mdx @@ -106,4 +106,4 @@ A `PUT` body sets, per budget, a limit or the unlimited flag; leave both out to { "token_limit": 2000000, "cost_unlimited": true, "note": "Research team" } ``` -`bucket` (default `all`) narrows a policy to `direct` traffic (chat without an agent) or `agent` traffic (anything that runs through an agent). A request must fit both its own bucket and `all`. The dashboard edits `all`. +`bucket` (default `all`) narrows a policy to `direct` traffic (chat without an agent) or `agent` traffic (anything that runs through an agent, whether or not the agent has an API key). A request must fit both its own bucket and `all`. The dashboard edits `all`. diff --git a/docsgpt/agents/headless_runner.py b/docsgpt/agents/headless_runner.py index 344e5fe3..c2a4d713 100644 --- a/docsgpt/agents/headless_runner.py +++ b/docsgpt/agents/headless_runner.py @@ -87,7 +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") + # An agent run is agent traffic whether or not the agent has a key yet. + is_agent_run = bool(agent_config.get("key") or _resolve_agent_id(agent_config)) + exceeded = QuotaService.check(owner, "agent" if is_agent_run else "direct") if exceeded is not None: raise QuotaExceededError(exceeded) diff --git a/docsgpt/api/answer/routes/answer.py b/docsgpt/api/answer/routes/answer.py index cdc63cca..b5819e3e 100644 --- a/docsgpt/api/answer/routes/answer.py +++ b/docsgpt/api/answer/routes/answer.py @@ -103,8 +103,8 @@ class AnswerResource(Resource, BaseAnswerResource): ) if not processor.decoded_token: return make_response({"error": "Unauthorized"}, 401) - if error := self.check_usage( - processor.agent_config, processor.decoded_token + if error := self.check_usage_on_resume( + processor, data["conversation_id"] ): return error stream = self.complete_stream( diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index c61cfd48..c2b3f929 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -211,6 +211,32 @@ class BaseAnswerResource: ) return None + def check_usage_on_resume(self, processor: Any, conversation_id: Any) -> Optional[Response]: + """Run ``check_usage`` for a tool continuation, releasing its claim on refusal. + + ``resume_from_tool_actions`` has already claimed the paused turn by the + time the limits can be checked (the agent config comes from the claimed + state). A refusal returns before ``complete_stream`` and its cleanup, so + the claim is released here; otherwise retries get a 409 until the stale + claim is reverted. + + Args: + processor: The ``StreamProcessor`` that resumed the turn. + conversation_id: The conversation whose pending state was claimed. + + Returns: + None, or the refusal Response. + """ + error = self.check_usage(processor.agent_config, processor.decoded_token) + if error is None or not conversation_id: + return error + user = processor.initial_user_id or (processor.decoded_token or {}).get("sub") + try: + ContinuationService().release_claim(str(conversation_id), user) + except Exception: + logger.exception("Failed to release resume claim after a usage refusal") + return error + def complete_stream( self, question: str, diff --git a/docsgpt/api/answer/routes/stream.py b/docsgpt/api/answer/routes/stream.py index fdc5a00b..54b62ec8 100644 --- a/docsgpt/api/answer/routes/stream.py +++ b/docsgpt/api/answer/routes/stream.py @@ -115,9 +115,9 @@ class StreamResource(Resource, BaseAnswerResource): status=401, mimetype="text/event-stream", ) - if error := self.check_usage( - processor.agent_config, processor.decoded_token - ): + if error := self.check_usage_on_resume( + processor, data["conversation_id"] + ): return error return Response( with_sse_keepalive( diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index cd208e84..a9902a4a 100644 --- a/docsgpt/api/v1/routes.py +++ b/docsgpt/api/v1/routes.py @@ -258,6 +258,8 @@ def chat_completions(): try: processor = StreamProcessor(internal_data, decoded_token) + # Set when this request took the resume claim, so a refusal can release it. + claimed_conversation_id = None if internal_data.get("tool_actions"): conversation_id = internal_data.get("conversation_id") @@ -282,6 +284,7 @@ def chat_completions(): claimed_state=pending_state, ) processor.conversation_id = conversation_id + claimed_conversation_id = conversation_id else: # Compatibility fallback for old/completed conversations and # clients that resend the full transcript without resumable @@ -338,7 +341,10 @@ def chat_completions(): ) helper = _V1AnswerHelper() - usage_error = helper.check_usage(processor.agent_config, processor.decoded_token) + if claimed_conversation_id: + usage_error = helper.check_usage_on_resume(processor, claimed_conversation_id) + else: + usage_error = helper.check_usage(processor.agent_config, processor.decoded_token) if usage_error: return usage_error diff --git a/docsgpt/core/models/openrouter.yaml b/docsgpt/core/models/openrouter.yaml index 6ac373ee..f0b2cffd 100644 --- a/docsgpt/core/models/openrouter.yaml +++ b/docsgpt/core/models/openrouter.yaml @@ -15,7 +15,7 @@ models: - id: deepseek/deepseek-v3.2 display_name: DeepSeek V3.2 - description: Open-weights reasoning model, very low cost (~$0.25 in / $0.38 out per 1M) + description: Open-weights reasoning model, very low cost (~$0.23 in / $0.34 out per 1M) context_window: 131072 attachments: [] supports_structured_output: true diff --git a/docsgpt/storage/db/repositories/token_usage.py b/docsgpt/storage/db/repositories/token_usage.py index 2c8038fa..fb256b55 100644 --- a/docsgpt/storage/db/repositories/token_usage.py +++ b/docsgpt/storage/db/repositories/token_usage.py @@ -144,16 +144,17 @@ class TokenUsageRepository: Args: user_id: The billable user (auth ``sub``). start: Inclusive window start. - bucket: ``all``, ``direct`` (rows without an agent key) or - ``agent`` (rows with one). + bucket: ``all``, ``agent`` (rows carrying an agent key or an agent + id) or ``direct`` (rows with neither). Rollup rows are excluded; side-channel calls count, they are real spend. """ clauses = ["user_id = :user_id", "timestamp >= :start", "source <> ALL(:rollup_sources)"] + # Keyless agents and workflow nodes carry an agent id without a key. if bucket == "direct": - clauses.append("api_key IS NULL") + clauses.append("api_key IS NULL AND agent_id IS NULL") elif bucket == "agent": - clauses.append("api_key IS NOT NULL") + clauses.append("(api_key IS NOT NULL OR agent_id IS NOT NULL)") elif bucket != "all": raise ValueError(f"unknown usage bucket: {bucket!r}") row = self._conn.execute( diff --git a/frontend/src/admin/UserQuotaModal.tsx b/frontend/src/admin/UserQuotaModal.tsx index f6ce58b2..241da310 100644 --- a/frontend/src/admin/UserQuotaModal.tsx +++ b/frontend/src/admin/UserQuotaModal.tsx @@ -1,4 +1,4 @@ -import { useCallback, useEffect, useState } from 'react'; +import { useCallback, useEffect, useRef, useState } from 'react'; import { useSelector } from 'react-redux'; import adminService from '../api/services/adminService'; @@ -26,15 +26,22 @@ export default function UserQuotaModal({ const [data, setData] = useState(null); const [teamNames, setTeamNames] = useState>({}); + // Bumped per request so a slow response for a previous user is discarded + // instead of showing (and letting the editor save) that user's policy. + const requestRef = useRef(0); + const load = useCallback(async () => { - if (!userId) return; + const request = ++requestRef.current; setData(null); + if (!userId) return; try { const [res, teamsJson] = await Promise.all([ adminService.getUserQuota(userId, token), teamsService.listAll(token).catch(() => ({})), ]); - setData(await res.json().catch(() => ({ success: false }))); + const json = await res.json().catch(() => ({ success: false })); + if (request !== requestRef.current) return; + setData(json); setTeamNames( Object.fromEntries( (teamsJson?.teams ?? []).map((team: any) => [ @@ -44,7 +51,7 @@ export default function UserQuotaModal({ ), ); } catch { - setData({ success: false }); + if (request === requestRef.current) setData({ success: false }); } }, [userId, token]); diff --git a/frontend/src/locale/de.json b/frontend/src/locale/de.json index 9c6879f7..df6dabfa 100644 --- a/frontend/src/locale/de.json +++ b/frontend/src/locale/de.json @@ -381,7 +381,11 @@ "resets": "Wird zurückgesetzt: {{resetsAt}}", "tokens": "Tokens", "cost": "Kosten", - "usedOf": "{{used}} von {{limit}}" + "usedOf": "{{used}} von {{limit}}", + "scope": { + "direct": "Chat ohne Agent", + "agent": "Über Agenten" + } } }, "logs": { diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index ba269327..cdf03a26 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -386,7 +386,11 @@ "resets": "Resets {{resetsAt}}", "tokens": "Tokens", "cost": "Cost", - "usedOf": "{{used}} of {{limit}}" + "usedOf": "{{used}} of {{limit}}", + "scope": { + "direct": "Chat without an agent", + "agent": "Through agents" + } } }, "logs": { diff --git a/frontend/src/locale/es.json b/frontend/src/locale/es.json index b844829b..76058628 100644 --- a/frontend/src/locale/es.json +++ b/frontend/src/locale/es.json @@ -381,7 +381,11 @@ "resets": "Se restablece el {{resetsAt}}", "tokens": "Tokens", "cost": "Coste", - "usedOf": "{{used}} de {{limit}}" + "usedOf": "{{used}} de {{limit}}", + "scope": { + "direct": "Chat sin agente", + "agent": "A través de agentes" + } } }, "logs": { diff --git a/frontend/src/locale/jp.json b/frontend/src/locale/jp.json index 7c5aec0b..18626d19 100644 --- a/frontend/src/locale/jp.json +++ b/frontend/src/locale/jp.json @@ -381,7 +381,11 @@ "resets": "{{resetsAt}} にリセット", "tokens": "トークン", "cost": "コスト", - "usedOf": "{{used}} / {{limit}}" + "usedOf": "{{used}} / {{limit}}", + "scope": { + "direct": "エージェントなしのチャット", + "agent": "エージェント経由" + } } }, "logs": { diff --git a/frontend/src/locale/ru.json b/frontend/src/locale/ru.json index 69c0c25f..081860c6 100644 --- a/frontend/src/locale/ru.json +++ b/frontend/src/locale/ru.json @@ -381,7 +381,11 @@ "resets": "Сброс: {{resetsAt}}", "tokens": "Токены", "cost": "Стоимость", - "usedOf": "{{used}} из {{limit}}" + "usedOf": "{{used}} из {{limit}}", + "scope": { + "direct": "Чат без агента", + "agent": "Через агентов" + } } }, "logs": { diff --git a/frontend/src/locale/zh-TW.json b/frontend/src/locale/zh-TW.json index 4de54e9d..7902cf5a 100644 --- a/frontend/src/locale/zh-TW.json +++ b/frontend/src/locale/zh-TW.json @@ -381,7 +381,11 @@ "resets": "{{resetsAt}} 重設", "tokens": "權杖", "cost": "費用", - "usedOf": "{{used}} / {{limit}}" + "usedOf": "{{used}} / {{limit}}", + "scope": { + "direct": "不使用代理的聊天", + "agent": "透過代理" + } } }, "logs": { diff --git a/frontend/src/locale/zh.json b/frontend/src/locale/zh.json index 6a542506..7e647c5e 100644 --- a/frontend/src/locale/zh.json +++ b/frontend/src/locale/zh.json @@ -381,7 +381,11 @@ "resets": "{{resetsAt}} 重置", "tokens": "令牌", "cost": "费用", - "usedOf": "{{used}} / {{limit}}" + "usedOf": "{{used}} / {{limit}}", + "scope": { + "direct": "不使用代理的聊天", + "agent": "通过代理" + } } }, "logs": { diff --git a/frontend/src/settings/Analytics.tsx b/frontend/src/settings/Analytics.tsx index 4eaccd5c..f67855bd 100644 --- a/frontend/src/settings/Analytics.tsx +++ b/frontend/src/settings/Analytics.tsx @@ -414,7 +414,7 @@ export default function Analytics({ agentId }: AnalyticsProps) {

{card.label}

diff --git a/frontend/src/settings/components/UsageQuota.tsx b/frontend/src/settings/components/UsageQuota.tsx index 7cb9ceff..30511081 100644 --- a/frontend/src/settings/components/UsageQuota.tsx +++ b/frontend/src/settings/components/UsageQuota.tsx @@ -66,7 +66,7 @@ function Meter({ export default function UsageQuota() { const { t, i18n } = useTranslation(); const token = useSelector(selectToken); - const [bucket, setBucket] = useState(null); + const [buckets, setBuckets] = useState([]); useEffect(() => { let cancelled = false; @@ -75,10 +75,7 @@ export default function UsageQuota() { .then((res: Response) => (res.ok ? res.json() : null)) .then((json: { buckets?: Bucket[] } | null) => { if (cancelled) return; - const buckets = json?.buckets ?? []; - setBucket( - buckets.find((b) => b.bucket === 'all') ?? buckets[0] ?? null, - ); + setBuckets(json?.buckets ?? []); }) .catch(() => undefined); return () => { @@ -86,14 +83,14 @@ export default function UsageQuota() { }; }, [token]); - if (!bucket) return null; + if (buckets.length === 0) return null; const number = new Intl.NumberFormat(i18n.language); const usd = new Intl.NumberFormat(i18n.language, { style: 'currency', currency: 'USD', }); - const reset = new Date(bucket.resets_at); + const reset = new Date(buckets[0].resets_at); const resetsAt = Number.isNaN(reset.getTime()) ? '' : new Intl.DateTimeFormat(i18n.language, { @@ -101,6 +98,12 @@ export default function UsageQuota() { timeStyle: 'short', }).format(reset); + // A request must fit its own bucket and ``all``, so each limited one is shown. + const scopeLabel = (name: string) => + name === 'direct' || name === 'agent' + ? t(`settings.analytics.quota.scope.${name}`) + : null; + return (

@@ -113,18 +116,27 @@ export default function UsageQuota() {

) : null}
-
- number.format(value)} - /> - usd.format(value)} - /> -
+ {buckets.map((bucket) => ( +
+ {scopeLabel(bucket.bucket) ? ( +

+ {scopeLabel(bucket.bucket)} +

+ ) : null} +
+ number.format(value)} + /> + usd.format(value)} + /> +
+
+ ))}
); } diff --git a/tests/quotas/test_enforcement.py b/tests/quotas/test_enforcement.py index b02dbae1..c6300b22 100644 --- a/tests/quotas/test_enforcement.py +++ b/tests/quotas/test_enforcement.py @@ -106,3 +106,70 @@ class TestHeadless: assert raised.value.exceeded.source == "instance" assert "10 of 10 tokens" in str(raised.value) + + + def test_a_keyless_agent_run_is_agent_traffic(self, db): + from docsgpt.agents.headless_runner import run_agent_headless + + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="owner", bucket="agent", token_limit=0) + config = {"user_id": "owner", "id": "22222222-2222-2222-2222-222222222222"} + + with patch("docsgpt.agents.headless_runner.RetrieverCreator"): + with pytest.raises(QuotaExceededError) as raised: + run_agent_headless(config, "hello") + + assert raised.value.exceeded.bucket == "agent" + + +class TestResumeRefusal: + def _processor(self, user_id="u1"): + from types import SimpleNamespace + + return SimpleNamespace(agent_config={}, decoded_token={"sub": user_id}, initial_user_id=user_id) + + def test_a_refused_resume_releases_its_claim(self, db, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=0) + with flask_app.app_context(), patch( + "docsgpt.api.answer.routes.base.ContinuationService" + ) as service: + response = BaseAnswerResource().check_usage_on_resume(self._processor(), "conv-1") + + assert response.status_code == 429 + service.return_value.release_claim.assert_called_once_with("conv-1", "u1") + + def test_an_admitted_resume_keeps_its_claim(self, db, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), patch( + "docsgpt.api.answer.routes.base.ContinuationService" + ) as service: + response = BaseAnswerResource().check_usage_on_resume(self._processor(), "conv-1") + + assert response is None + service.return_value.release_claim.assert_not_called() + + def test_no_claim_means_nothing_to_release(self, db, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=0) + with flask_app.app_context(), patch( + "docsgpt.api.answer.routes.base.ContinuationService" + ) as service: + response = BaseAnswerResource().check_usage_on_resume(self._processor(), None) + + assert response.status_code == 429 + service.return_value.release_claim.assert_not_called() + + def test_a_failed_release_still_returns_the_refusal(self, db, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=0) + with flask_app.app_context(), patch( + "docsgpt.api.answer.routes.base.ContinuationService" + ) as service: + service.return_value.release_claim.side_effect = RuntimeError("db down") + response = BaseAnswerResource().check_usage_on_resume(self._processor(), "conv-1") + + assert response.status_code == 429 diff --git a/tests/storage/db/repositories/test_token_usage.py b/tests/storage/db/repositories/test_token_usage.py index 9698c610..6ce4cb96 100644 --- a/tests/storage/db/repositories/test_token_usage.py +++ b/tests/storage/db/repositories/test_token_usage.py @@ -61,6 +61,14 @@ class TestUsageTotals: repo.insert(user_id="u-tot", prompt_tokens=100, generated_tokens=10, cost=0.5) repo.insert(user_id="u-tot", api_key="k", prompt_tokens=20, generated_tokens=2, cost=0.25) repo.insert(user_id="u-tot", prompt_tokens=7, generated_tokens=0, cost=0.125, source="title") + # A keyless agent (or workflow node): an agent id without a key. + from docsgpt.storage.db.repositories.agents import AgentsRepository + + agent = AgentsRepository(repo._conn).create("u-tot", "keyless", "draft") + repo.insert( + user_id="u-tot", agent_id=str(agent["id"]), + prompt_tokens=4, generated_tokens=0, cost=0.0625, + ) repo.insert(user_id="u-tot", prompt_tokens=999, generated_tokens=0, source="schedule") repo.insert(user_id="u-other", prompt_tokens=999, generated_tokens=0, cost=9) repo.insert( @@ -70,7 +78,7 @@ class TestUsageTotals: @pytest.mark.parametrize( "bucket, expected", - [("all", (139, 0.875)), ("direct", (117, 0.625)), ("agent", (22, 0.25))], + [("all", (143, 0.9375)), ("direct", (117, 0.625)), ("agent", (26, 0.3125))], ) def test_totals_per_bucket(self, pg_conn, bucket, expected): repo = _repo(pg_conn) From 3a30f5cce9d57464d3c0bfb61ee1de2cab9a3d2d Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 13:36:22 +0100 Subject: [PATCH 12/14] feat(pricing): rates for the default DocsGPT model $0.15 input, $0.50 output and $0.03 cached input per 1M tokens, so cost budgets see usage of the default model instead of recording it at $0. --- docsgpt/core/models/docsgpt.yaml | 3 +++ tests/test_pricing.py | 11 ++++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/docsgpt/core/models/docsgpt.yaml b/docsgpt/core/models/docsgpt.yaml index b65b4fbf..4f494434 100644 --- a/docsgpt/core/models/docsgpt.yaml +++ b/docsgpt/core/models/docsgpt.yaml @@ -7,3 +7,6 @@ models: supports_tools: true attachments: [image] context_window: 1048576 + input_cost_per_million: 0.15 + output_cost_per_million: 0.5 + cached_input_cost_per_million: 0.03 diff --git a/tests/test_pricing.py b/tests/test_pricing.py index 81aa324f..d8a6ca69 100644 --- a/tests/test_pricing.py +++ b/tests/test_pricing.py @@ -118,7 +118,7 @@ class TestCatalogFields: self._load(tmp_path, "provider: openai\nmodels:\n - id: m\n input_cost_per_million: -1\n") def test_hosted_builtin_models_are_priced(self): - hosted = {"anthropic", "deepseek", "google", "groq", "novita", "openai", "openrouter"} + hosted = {"anthropic", "deepseek", "docsgpt", "google", "groq", "novita", "openai", "openrouter"} catalogs = [ c for c in load_model_yamls([BUILTIN_MODELS_DIR]) if c.source_path.stem in hosted ] @@ -128,3 +128,12 @@ class TestCatalogFields: caps = model.capabilities assert caps.input_cost_per_million is not None, model.id assert caps.output_cost_per_million is not None, model.id + + def test_default_docsgpt_model_rates(self): + (model,) = [ + m for c in load_model_yamls([BUILTIN_MODELS_DIR]) for m in c.models if m.id == "docsgpt-local" + ] + caps = model.capabilities + assert (caps.input_cost_per_million, caps.output_cost_per_million) == (0.15, 0.5) + assert caps.cached_input_cost_per_million == 0.03 + assert caps.cache_write_cost_per_million is None From b5296df8a934e8fec7cd3f0b2aa4b341ec8eab17 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 14:30:48 +0100 Subject: [PATCH 13/14] feat(pricing): cached-input rates for gpt-5.4-mini and gpt-5.4-nano Checked against OpenAI's pricing page: gpt-5.5 at $5 / $30 (cached $0.50) was already right. The mini and nano models declared no cached rate, so cached prompt tokens were billed at the full input rate. --- docsgpt/core/models/openai.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docsgpt/core/models/openai.yaml b/docsgpt/core/models/openai.yaml index 598e7b3c..a4b73a3a 100644 --- a/docsgpt/core/models/openai.yaml +++ b/docsgpt/core/models/openai.yaml @@ -12,6 +12,7 @@ models: context_window: 1050000 api_flavor: responses reasoning_effort: medium + # Short-context rates. Prompts over 272K tokens bill at $10 / $45 (cached $1). input_cost_per_million: 5.0 output_cost_per_million: 30.0 cached_input_cost_per_million: 0.5 @@ -20,8 +21,10 @@ models: description: Cost-efficient GPT-5.4-class model for high-volume coding, computer use, and subagent workloads input_cost_per_million: 0.75 output_cost_per_million: 4.5 + cached_input_cost_per_million: 0.075 - id: gpt-5.4-nano display_name: GPT-5.4 Nano description: Cheapest GPT-5.4-class model, optimized for simple high-volume tasks where speed and cost matter most input_cost_per_million: 0.2 output_cost_per_million: 1.25 + cached_input_cost_per_million: 0.02 From 1e14605ee7c18636ba46e5f7bd158b1ef0d91398 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 15:26:29 +0100 Subject: [PATCH 14/14] fix(quotas): keyless agent chat bucket, keep disabled policies disabled, cached rates - check_usage treats a request through a keyless (draft) agent as agent traffic, matching how its usage rows are bucketed and the headless rule - dashboard edits carry the stored enabled flag instead of re-enabling the policy; disabled policies are labelled in the Quotas tab and the editor - quota 429s send x-should-retry: false so OpenAI SDK clients do not retry a refusal that cannot succeed before the reset - cached-input and cache-write rates for Anthropic, OpenRouter and Groq gpt-oss-120b; refresh OpenRouter deepseek-v3.2 list prices - UsageQuota reuses usagePercent; docs note that a user override needs an existing user --- docs/content/Deploying/Usage-Quotas.mdx | 2 +- docsgpt/api/answer/routes/answer.py | 4 ++- docsgpt/api/answer/routes/base.py | 13 +++++++--- docsgpt/api/answer/routes/stream.py | 4 ++- docsgpt/api/v1/routes.py | 4 ++- docsgpt/core/models/anthropic.yaml | 6 +++++ docsgpt/core/models/groq.yaml | 1 + docsgpt/core/models/openrouter.yaml | 9 ++++--- docsgpt/quotas/http.py | 2 ++ frontend/src/admin/QuotaEditor.tsx | 7 +++++- frontend/src/admin/Quotas.tsx | 10 +++++--- frontend/src/admin/quotaUtils.test.ts | 8 ++++++ frontend/src/admin/quotaUtils.ts | 8 +++++- .../src/settings/components/UsageQuota.tsx | 6 ++--- tests/quotas/test_enforcement.py | 25 ++++++++++++++++--- 15 files changed, 87 insertions(+), 22 deletions(-) diff --git a/docs/content/Deploying/Usage-Quotas.mdx b/docs/content/Deploying/Usage-Quotas.mdx index c7391625..e4ffbaf1 100644 --- a/docs/content/Deploying/Usage-Quotas.mdx +++ b/docs/content/Deploying/Usage-Quotas.mdx @@ -97,7 +97,7 @@ Every admin endpoint requires the admin role, and every change is written to the | `GET` | `/api/admin/quotas` | All policies by layer, the current window, and used models without a price. | | `PUT` `DELETE` | `/api/admin/quotas/instance` | The instance default. | | `GET` `PUT` `DELETE` | `/api/admin/quotas/teams/` | A team's per-member allowance. | -| `GET` `PUT` `DELETE` | `/api/admin/quotas/users/` | A user's override. `GET` also returns the limits the user ends up with, the layer each came from, and their usage. | +| `GET` `PUT` `DELETE` | `/api/admin/quotas/users/` | A user's override; the user must already exist (SCIM-provisioned, or signed in once), otherwise `404`. `GET` also returns the limits the user ends up with, the layer each came from, and their usage. | | `GET` | `/api/user/quota` | The caller's own limits, usage and reset time. | A `PUT` body sets, per budget, a limit or the unlimited flag; leave both out to defer to the next layer: diff --git a/docsgpt/api/answer/routes/answer.py b/docsgpt/api/answer/routes/answer.py index b5819e3e..7cdcd722 100644 --- a/docsgpt/api/answer/routes/answer.py +++ b/docsgpt/api/answer/routes/answer.py @@ -132,7 +132,9 @@ class AnswerResource(Resource, BaseAnswerResource): return make_response({"error": "Unauthorized"}, 401) if error := self.check_usage( - processor.agent_config, processor.decoded_token + processor.agent_config, + processor.decoded_token, + agent_id=processor.agent_id, ): return error diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index c2b3f929..e0cfe4cd 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -115,7 +115,10 @@ class BaseAnswerResource: return prepared def check_usage( - self, agent_config: Dict, decoded_token: Optional[Dict] = None + self, + agent_config: Dict, + decoded_token: Optional[Dict] = None, + agent_id: Optional[str] = None, ) -> Optional[Response]: """Refuse the request when a usage limit is exhausted. @@ -127,6 +130,8 @@ class BaseAnswerResource: agent_config: The config dict of agent instance decoded_token: The request's resolved identity; its ``sub`` is the billable user. + agent_id: The agent the request runs through. A draft agent has no + key, but its usage rows carry the agent id, so it is agent traffic. Returns: None or Response if either of limits exceeded. @@ -134,7 +139,7 @@ class BaseAnswerResource: """ 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") + exceeded = QuotaService.check(user_id, "agent" if api_key or agent_id else "direct") if exceeded is not None: return quota_exceeded_response(exceeded) if not api_key: @@ -227,7 +232,9 @@ class BaseAnswerResource: Returns: None, or the refusal Response. """ - error = self.check_usage(processor.agent_config, processor.decoded_token) + error = self.check_usage( + processor.agent_config, processor.decoded_token, agent_id=processor.agent_id + ) if error is None or not conversation_id: return error user = processor.initial_user_id or (processor.decoded_token or {}).get("sub") diff --git a/docsgpt/api/answer/routes/stream.py b/docsgpt/api/answer/routes/stream.py index 54b62ec8..cdee274b 100644 --- a/docsgpt/api/answer/routes/stream.py +++ b/docsgpt/api/answer/routes/stream.py @@ -154,7 +154,9 @@ class StreamResource(Resource, BaseAnswerResource): ) if error := self.check_usage( - processor.agent_config, processor.decoded_token + processor.agent_config, + processor.decoded_token, + agent_id=processor.agent_id, ): return error should_persist, visibility = resolve_persistence( diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index a9902a4a..fa915d24 100644 --- a/docsgpt/api/v1/routes.py +++ b/docsgpt/api/v1/routes.py @@ -344,7 +344,9 @@ def chat_completions(): if claimed_conversation_id: usage_error = helper.check_usage_on_resume(processor, claimed_conversation_id) else: - usage_error = helper.check_usage(processor.agent_config, processor.decoded_token) + usage_error = helper.check_usage( + processor.agent_config, processor.decoded_token, agent_id=processor.agent_id + ) if usage_error: return usage_error diff --git a/docsgpt/core/models/anthropic.yaml b/docsgpt/core/models/anthropic.yaml index 784ab8e3..34a2cedc 100644 --- a/docsgpt/core/models/anthropic.yaml +++ b/docsgpt/core/models/anthropic.yaml @@ -12,6 +12,8 @@ models: supports_structured_output: true input_cost_per_million: 5.0 output_cost_per_million: 25.0 + cached_input_cost_per_million: 0.5 + cache_write_cost_per_million: 6.25 - id: claude-sonnet-4-6 display_name: Claude Sonnet 4.6 @@ -20,6 +22,8 @@ models: supports_structured_output: true input_cost_per_million: 3.0 output_cost_per_million: 15.0 + cached_input_cost_per_million: 0.3 + cache_write_cost_per_million: 3.75 - id: claude-haiku-4-5 display_name: Claude Haiku 4.5 @@ -27,3 +31,5 @@ models: supports_structured_output: true input_cost_per_million: 1.0 output_cost_per_million: 5.0 + cached_input_cost_per_million: 0.1 + cache_write_cost_per_million: 1.25 diff --git a/docsgpt/core/models/groq.yaml b/docsgpt/core/models/groq.yaml index c6e28d7a..a4d7edfd 100644 --- a/docsgpt/core/models/groq.yaml +++ b/docsgpt/core/models/groq.yaml @@ -10,6 +10,7 @@ models: supports_structured_output: true input_cost_per_million: 0.15 output_cost_per_million: 0.6 + cached_input_cost_per_million: 0.075 - id: llama-3.3-70b-versatile display_name: Llama 3.3 70B Versatile description: Meta's Llama 3.3 70B for general-purpose chat with parallel tool use diff --git a/docsgpt/core/models/openrouter.yaml b/docsgpt/core/models/openrouter.yaml index f0b2cffd..2957fa98 100644 --- a/docsgpt/core/models/openrouter.yaml +++ b/docsgpt/core/models/openrouter.yaml @@ -15,12 +15,13 @@ models: - id: deepseek/deepseek-v3.2 display_name: DeepSeek V3.2 - description: Open-weights reasoning model, very low cost (~$0.23 in / $0.34 out per 1M) + description: Open-weights reasoning model, very low cost (~$0.27 in / $0.40 out per 1M) context_window: 131072 attachments: [] supports_structured_output: true - input_cost_per_million: 0.23 - output_cost_per_million: 0.34 + input_cost_per_million: 0.269 + output_cost_per_million: 0.4 + cached_input_cost_per_million: 0.1345 - id: anthropic/claude-sonnet-4.6 display_name: Claude Sonnet 4.6 (via OpenRouter) @@ -29,3 +30,5 @@ models: supports_structured_output: true input_cost_per_million: 3.0 output_cost_per_million: 15.0 + cached_input_cost_per_million: 0.3 + cache_write_cost_per_million: 3.75 diff --git a/docsgpt/quotas/http.py b/docsgpt/quotas/http.py index 494b87cf..c9a1a180 100644 --- a/docsgpt/quotas/http.py +++ b/docsgpt/quotas/http.py @@ -11,4 +11,6 @@ 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) + # The reset can be weeks away; the OpenAI SDKs would otherwise retry with backoff. + response.headers["x-should-retry"] = "false" return response diff --git a/frontend/src/admin/QuotaEditor.tsx b/frontend/src/admin/QuotaEditor.tsx index e962ed46..7b9de108 100644 --- a/frontend/src/admin/QuotaEditor.tsx +++ b/frontend/src/admin/QuotaEditor.tsx @@ -185,7 +185,7 @@ export default function QuotaEditor({ if (policy) remove(); return; } - const result = formToPolicy(form); + const result = formToPolicy(form, policy); if (!result.ok) { setError(result.error); return; @@ -196,6 +196,11 @@ export default function QuotaEditor({ return (

{inheritHint}

+ {policy && !policy.enabled ? ( +

+ This policy is disabled and is not enforced. Saving keeps it disabled. +

+ ) : null} ; if (!data?.success) return ; - const bucketPill = (policy: QuotaPolicy) => - isAll(policy) ? null : {policy.bucket} traffic; + const bucketPill = (policy: QuotaPolicy) => ( + <> + {isAll(policy) ? null : {policy.bucket} traffic} + {policy.enabled ? null : Disabled} + + ); return (
@@ -163,7 +167,7 @@ export default function Quotas() {

{instancePolicy - ? `Tokens: ${describeBudget(instancePolicy.token_limit, instancePolicy.token_unlimited, 'tokens')} · Cost: ${describeBudget(instancePolicy.cost_limit_usd, instancePolicy.cost_unlimited, 'cost')}` + ? `${instancePolicy.enabled ? '' : 'Disabled · '}Tokens: ${describeBudget(instancePolicy.token_limit, instancePolicy.token_unlimited, 'tokens')} · Cost: ${describeBudget(instancePolicy.cost_limit_usd, instancePolicy.cost_unlimited, 'cost')}` : 'No default: users without a team allowance or override are unlimited.'}

diff --git a/frontend/src/admin/quotaUtils.test.ts b/frontend/src/admin/quotaUtils.test.ts index 60e4f468..2afad959 100644 --- a/frontend/src/admin/quotaUtils.test.ts +++ b/frontend/src/admin/quotaUtils.test.ts @@ -54,6 +54,7 @@ describe('formToPolicy', () => { ok: true, policy: { bucket: 'all', + enabled: true, token_limit: 5000, token_unlimited: false, cost_limit_usd: 2.5, @@ -63,6 +64,13 @@ describe('formToPolicy', () => { }); }); + it('keeps a disabled policy disabled', () => { + const form = { ...base, tokenMode: 'limit' as const, tokenLimit: '10' }; + const stored = policy({ enabled: false }); + const result = formToPolicy(form, stored); + expect(result.ok && result.policy.enabled).toBe(false); + }); + it('sends unlimited without a limit', () => { const result = formToPolicy({ ...base, diff --git a/frontend/src/admin/quotaUtils.ts b/frontend/src/admin/quotaUtils.ts index 10038574..02bc2f6b 100644 --- a/frontend/src/admin/quotaUtils.ts +++ b/frontend/src/admin/quotaUtils.ts @@ -66,9 +66,15 @@ export function isEmptyForm(form: QuotaForm): boolean { return form.tokenMode === 'inherit' && form.costMode === 'inherit'; } -export function formToPolicy(form: QuotaForm): FormResult { +// ``existing`` carries the stored ``enabled`` flag through an edit: the form has +// no control for it, and a body without it would switch the policy back on. +export function formToPolicy( + form: QuotaForm, + existing?: QuotaPolicy | null, +): FormResult { const policy: Record = { bucket: 'all', + enabled: existing?.enabled ?? true, token_limit: null, token_unlimited: form.tokenMode === 'unlimited', cost_limit_usd: null, diff --git a/frontend/src/settings/components/UsageQuota.tsx b/frontend/src/settings/components/UsageQuota.tsx index 30511081..a2147beb 100644 --- a/frontend/src/settings/components/UsageQuota.tsx +++ b/frontend/src/settings/components/UsageQuota.tsx @@ -2,6 +2,7 @@ import { useEffect, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; +import { usagePercent } from '../../admin/quotaUtils'; import userService from '../../api/services/userService'; import { selectToken } from '../../preferences/preferenceSlice'; @@ -24,10 +25,7 @@ function Meter({ }) { const { t } = useTranslation(); if (budget.limit === null) return null; - const percent = - budget.limit <= 0 - ? 100 - : Math.min(100, Math.max(0, (budget.used / budget.limit) * 100)); + const percent = usagePercent(budget.used, budget.limit); const tone = percent >= 100 ? 'bg-red-500' diff --git a/tests/quotas/test_enforcement.py b/tests/quotas/test_enforcement.py index c6300b22..0a54974c 100644 --- a/tests/quotas/test_enforcement.py +++ b/tests/quotas/test_enforcement.py @@ -30,11 +30,11 @@ def _spend(conn, user_id, tokens, api_key=None): TokenUsageRepository(conn).insert(user_id=user_id, api_key=api_key, prompt_tokens=tokens) -def _check(flask_app, agent_config, decoded_token=None): +def _check(flask_app, agent_config, decoded_token=None, agent_id=None): from docsgpt.api.answer.routes.base import BaseAnswerResource with flask_app.app_context(): - return BaseAnswerResource().check_usage(agent_config, decoded_token) + return BaseAnswerResource().check_usage(agent_config, decoded_token, agent_id=agent_id) class TestCheckUsage: @@ -90,6 +90,23 @@ class TestCheckUsage: _spend(db, "u1", 500) assert _check(flask_app, {}, {"sub": "u1"}) is None + def test_a_keyless_agent_chat_is_agent_traffic(self, db, flask_app): + agent_id = str(AgentsRepository(db).create("u1", "draft", "draft")["id"]) + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", bucket="agent", token_limit=10) + TokenUsageRepository(db).insert(user_id="u1", agent_id=agent_id, prompt_tokens=10) + + response = _check(flask_app, {"user_api_key": None}, {"sub": "u1"}, agent_id=agent_id) + + assert response.status_code == 429 + assert json.loads(response.data)["bucket"] == "agent" + # The same spend leaves chat without an agent alone. + assert _check(flask_app, {}, {"sub": "u1"}) is None + + def test_a_refusal_tells_sdk_clients_not_to_retry(self, db, flask_app): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=0) + response = _check(flask_app, {}, {"sub": "u1"}) + assert response.headers["x-should-retry"] == "false" + class TestHeadless: def test_exhausted_owner_is_refused_before_the_run(self, db): @@ -125,7 +142,9 @@ class TestResumeRefusal: def _processor(self, user_id="u1"): from types import SimpleNamespace - return SimpleNamespace(agent_config={}, decoded_token={"sub": user_id}, initial_user_id=user_id) + return SimpleNamespace( + agent_config={}, decoded_token={"sub": user_id}, initial_user_id=user_id, agent_id=None + ) def test_a_refused_resume_releases_its_claim(self, db, flask_app): from docsgpt.api.answer.routes.base import BaseAnswerResource