mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-11 03:12:55 +00:00
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.
This commit is contained in:
1 parent
69f55b74cb
commit
4bd259fc09
20 files changed
+202
-45
No files matched your search
@@ -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`.
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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<any | null>(null);
|
||||
const [teamNames, setTeamNames] = useState<Record<string, string>>({});
|
||||
|
||||
// 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]);
|
||||
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -381,7 +381,11 @@
|
||||
"resets": "{{resetsAt}} にリセット",
|
||||
"tokens": "トークン",
|
||||
"cost": "コスト",
|
||||
"usedOf": "{{used}} / {{limit}}"
|
||||
"usedOf": "{{used}} / {{limit}}",
|
||||
"scope": {
|
||||
"direct": "エージェントなしのチャット",
|
||||
"agent": "エージェント経由"
|
||||
}
|
||||
}
|
||||
},
|
||||
"logs": {
|
||||
|
||||
@@ -381,7 +381,11 @@
|
||||
"resets": "Сброс: {{resetsAt}}",
|
||||
"tokens": "Токены",
|
||||
"cost": "Стоимость",
|
||||
"usedOf": "{{used}} из {{limit}}"
|
||||
"usedOf": "{{used}} из {{limit}}",
|
||||
"scope": {
|
||||
"direct": "Чат без агента",
|
||||
"agent": "Через агентов"
|
||||
}
|
||||
}
|
||||
},
|
||||
"logs": {
|
||||
|
||||
@@ -381,7 +381,11 @@
|
||||
"resets": "{{resetsAt}} 重設",
|
||||
"tokens": "權杖",
|
||||
"cost": "費用",
|
||||
"usedOf": "{{used}} / {{limit}}"
|
||||
"usedOf": "{{used}} / {{limit}}",
|
||||
"scope": {
|
||||
"direct": "不使用代理的聊天",
|
||||
"agent": "透過代理"
|
||||
}
|
||||
}
|
||||
},
|
||||
"logs": {
|
||||
|
||||
@@ -381,7 +381,11 @@
|
||||
"resets": "{{resetsAt}} 重置",
|
||||
"tokens": "令牌",
|
||||
"cost": "费用",
|
||||
"usedOf": "{{used}} / {{limit}}"
|
||||
"usedOf": "{{used}} / {{limit}}",
|
||||
"scope": {
|
||||
"direct": "不使用代理的聊天",
|
||||
"agent": "通过代理"
|
||||
}
|
||||
}
|
||||
},
|
||||
"logs": {
|
||||
|
||||
@@ -414,7 +414,7 @@ export default function Analytics({ agentId }: AnalyticsProps) {
|
||||
<div
|
||||
key={card.label}
|
||||
title={card.hint}
|
||||
className={`border-border dark:border-border rounded-2xl border px-6 py-5${card.hint ? 'cursor-help' : ''}`}
|
||||
className={`border-border dark:border-border rounded-2xl border px-6 py-5${card.hint ? ' cursor-help' : ''}`}
|
||||
>
|
||||
<p className="text-muted-foreground text-sm">{card.label}</p>
|
||||
<p className="text-foreground dark:text-foreground mt-1 text-2xl font-bold">
|
||||
|
||||
@@ -66,7 +66,7 @@ function Meter({
|
||||
export default function UsageQuota() {
|
||||
const { t, i18n } = useTranslation();
|
||||
const token = useSelector(selectToken);
|
||||
const [bucket, setBucket] = useState<Bucket | null>(null);
|
||||
const [buckets, setBuckets] = useState<Bucket[]>([]);
|
||||
|
||||
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 (
|
||||
<div className="border-border mb-6 rounded-2xl border px-6 py-5">
|
||||
<div className="flex flex-wrap items-baseline justify-between gap-2">
|
||||
@@ -113,18 +116,27 @@ export default function UsageQuota() {
|
||||
</p>
|
||||
) : null}
|
||||
</div>
|
||||
<div className="mt-3 flex flex-wrap gap-6">
|
||||
<Meter
|
||||
label={t('settings.analytics.quota.tokens')}
|
||||
budget={bucket.tokens}
|
||||
format={(value) => number.format(value)}
|
||||
/>
|
||||
<Meter
|
||||
label={t('settings.analytics.quota.cost')}
|
||||
budget={bucket.cost}
|
||||
format={(value) => usd.format(value)}
|
||||
/>
|
||||
</div>
|
||||
{buckets.map((bucket) => (
|
||||
<div key={bucket.bucket} className="mt-3">
|
||||
{scopeLabel(bucket.bucket) ? (
|
||||
<p className="text-muted-foreground mb-1 text-xs">
|
||||
{scopeLabel(bucket.bucket)}
|
||||
</p>
|
||||
) : null}
|
||||
<div className="flex flex-wrap gap-6">
|
||||
<Meter
|
||||
label={t('settings.analytics.quota.tokens')}
|
||||
budget={bucket.tokens}
|
||||
format={(value) => number.format(value)}
|
||||
/>
|
||||
<Meter
|
||||
label={t('settings.analytics.quota.cost')}
|
||||
budget={bucket.cost}
|
||||
format={(value) => usd.format(value)}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user