From 4bd259fc091932a6ee0c017bb3595111089b5ea6 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 12:44:26 +0100 Subject: [PATCH] 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)