mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 10:14:32 +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
@@ -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(
|
||||
|
||||
Reference in new issue
Block a user