Skip to content

fix: chat user token quota #7814

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wxg0103 merged 1 commit into from
Nov 10, 2026
Merged

fix: chat user token quota #7814

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 36 additions & 25 deletions apps/chat/serializers/chat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,32 +351,42 @@ def on_complete(wf_manage, error):
chat_record_id_str,
FailureContent(str(uuid_utils.uuid7()), str(error), Status.SUCCESS, None, None).to_dict(),
)
messages = aggregation.get_contents()
old_details = None
chat_record = None
if chat_record_id is not None:
chat_record = QuerySet(ChatRecord).filter(id=chat_record_id).first()
if chat_record:
old_details = chat_record.details
if position and chat_record.messages:
messages = list({m.get("id"): m for m in [*chat_record.messages, *messages]}.values())
details = wf_manage.get_details(position=position, old_details=old_details)
run_time = time.time() - wf_manage.start_time
self.update_chat_record(chat_user_id, chat_record_id_str, wf_manage.context, messages, details, run_time)
self.update_chat_variable(chat_id, chat_variable, wf_manage.context.get("chat"))
ChatCountSerializer(data={"chat_id": chat_id}).update_chat()
# 表单续跑时 message_dict.content 为空;保留原记录里的用户问题,避免 WORKFLOW 历史丢问题
question = chat_record.question if (chat_record and chat_record.question) else message_dict
ChatHistory(chat_id).append(
ChatRecord(
id=chat_record_id_str,
chat_id=chat_id,
question=question,
messages=messages,
details=details,
create_time=timezone.now(),
try:
messages = aggregation.get_contents()
old_details = None
chat_record = None
if chat_record_id is not None:
chat_record = QuerySet(ChatRecord).filter(id=chat_record_id).first()
if chat_record:
old_details = chat_record.details
if position and chat_record.messages:
messages = list({m.get("id"): m for m in [*chat_record.messages, *messages]}.values())
details = wf_manage.get_details(position=position, old_details=old_details)
run_time = time.time() - wf_manage.start_time
self.update_chat_record(
chat_user_id, chat_record_id_str, wf_manage.context, messages, details, run_time
)
)
self.update_chat_variable(chat_id, chat_variable, wf_manage.context.get("chat"))
ChatCountSerializer(data={"chat_id": chat_id}).update_chat()
# 表单续跑时 message_dict.content 为空;保留原记录里的用户问题,避免 WORKFLOW 历史丢问题
question = chat_record.question if (chat_record and chat_record.question) else message_dict
ChatHistory(chat_id).append(
ChatRecord(
id=chat_record_id_str,
chat_id=chat_id,
question=question,
messages=messages,
details=details,
create_time=timezone.now(),
)
)
except Exception as e:
if not error:
result_queue.put(("error", e))
message_queue.produce(
chat_record_id_str,
FailureContent(str(uuid_utils.uuid7()), str(e), Status.SUCCESS, None, None).to_dict(),
)
result_queue.put(("done", None))
message_queue.produce_done(chat_record_id_str)

Expand Down Expand Up @@ -454,6 +464,7 @@ def chat(self, instance: dict, base_to_response: BaseToResponse = SystemToRespon
self.is_valid_chat()
application = self.get_application()
self.is_valid_intraday_access_num()
ChatUserTokenQuota.check_available(self.data.get("chat_user_id"))
return self.chat_work_flow(application, instance, base_to_response)


@shaohuzhang1 shaohuzhang1 Nov 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is the review of the provided code diff, covering potential issues, logic errors, and optimization suggestions.

1. Critical Logic Bug in Error Handling (Status Mismatch)

Location: In the new except block within on_complete:

message_queue.produce(
    chat_record_id_str,
    FailureContent(str(uuid_utils.uuid7()), str(e), Status.SUCCESS, None, None).to_dict(),
)

Issue: You are sending a Failure content (FailureContent) but setting the status to Status.SUCCESS.

  • This will likely cause client-side confusion or inconsistent state, as the client receives a payload marked as "Success" containing error information.
  • It contradicts the behavior seen earlier in the file where FailureContent is likely paired with a failure status (or at least a consistent status). If this block is meant to handle exceptions during post-processing, the status should reflect the failure (e.g., Status.FAILURE).

Suggestion: Check what FailureContent expects for status. Usually, it should be a failure state.

# Hypothetical correction
FailureContent(..., Status.FAILURE, ...)

2. Resource Leakage / Duplicate Notification

Location: Inside the try...except block.
Issue: If an exception occurs after the database updates (update_chat_record, etc.) but before result_queue.put(("done", None)), the except block triggers.

  • However, notice that after the try/except block, the code unconditionally executes: 
    result_queue.put(("done", None))
message_queue.produce_done(chat_record_id_str)
  • This means even if an error occurred (except block ran), the system still tells the client "Done" via produce_done. The client might receive both the error message and then a "done" signal, depending on how queues are consumed.
  • If the except block sends a failure signal to the message queue, does produce_done overwrite it or send another message? Typically, produce_done signifies normal termination. If an error occurred, you should probably raise the exception after logging/sending error info, or structure the flow so produce_done is only called on success.

Suggestion: If the exception is handled by sending an error message, you should stop execution. Add a return statement at the end of the except block or re-raise the error.

            except Exception as e:
                if not error:
                    result_queue.put(("error", e))
                    message_queue.produce(...)
                # Consider re-raising if outer handlers need to know about failure
                # Or ensure produce_done isn't reached if we want to halt strictly on error
                raise

(Note: If on_complete is a callback thread function, re-raising might kill the thread silently. Ensure the error propagation strategy matches your architecture.)

3. Security & Validation: Token Quota Check

Location: ChatUserTokenQuota.check_available(self.data.get("chat_user_id")) in chat(self, ...) method.
Issue:

  • NPE Risk: self.data.get("chat_user_id") could return None if the key is missing. Depending on check_available implementation, this might throw an AttributeError (if it tries .id or similar on None) or silently pass (if it treats None as valid).
  • Validation Order: You call is_valid_chat() before checking the quota. If is_valid_chat validates user existence, fetching self.data["chat_user_id"] directly (assuming it exists) or retrieving the validated user object would be safer than relying on dictionary access which might be untrusted input.

Suggestion:

# Safer approach if is_valid_chat already ensures user validity:
user_id = self.data["chat_user_id"] # Let KeyError fail fast if logic is broken
ChatUserTokenQuota.check_available(user_id)

4. Performance: N+1 Query / Redundant DB Fetch

Location: on_complete function.
Issue:

if chat_record_id is not None:
    chat_record = QuerySet(ChatRecord).filter(id=chat_record_id).first()

You fetch the ChatRecord from the database again inside on_complete, even though this workflow likely started with this record.

  • If on_complete is triggered asynchronously or in a separate context, this is necessary.
  • However, look at wf_manage.start_time. If wf_manage holds references to original records, consider passing the old details or record ID more efficiently to avoid a fresh DB round-trip in the callback path, which is critical for latency in high-throughput systems.
  • Minor: Using Model.objects.filter(id=id).first() is fine, but if chat_record_id is a primary key, ensure you're not bypassing caching layers or using inefficient ORMs.

5. Data Consistency: Message Deduplication

Location:

messages = list({m.get("id"): m for m in [*chat_record.messages, *messages]}.values())

Issue:

  • Missing IDs: If any message in chat_record.messages or the new messages does not have an "id" key, m.get("id") returns None.
  • Collision: All messages without IDs will collide on the key None in the dictionary comprehension, effectively discarding all but one random message that lacks an ID.
  • Ordering: Dictionaries preserve order in Python 3.7+, but using dict comprehensions for deduplication relies on last-write-wins. Ensure this is the desired behavior (newest message overrides old).

Suggestion: Verify that all messages always have unique IDs. If not, fallback to UUID generation for missing ones to prevent data loss.

6. Code Style & Readability

  • Broad Exception Catching: except Exception as e is very broad. It catches things like Ctrl+C signals (though less likely in threads) or SystemExit. Consider catching specific exceptions expected from DB operations or serialization errors.
  • Magic Status: As mentioned in point build(deps): bump the pip group across 1 directory with 2 updates #1, Status.SUCCESS inside a FailureContent container is highly suspicious. This is a code smell indicating a copy-paste error.

Summary of Recommended Actions

  1. Fix Status: Change Status.SUCCESS to appropriate failure status in the except block's FailureContent.
  2. Control Flow: Ensure produce_done isn't sent after an unrecovered error in on_complete, or clarify the protocol.
  3. Null Safety: Handle potential None in ChatUserTokenQuota.check_available.
  4. Robust Deduplication: Protect the dict-comprehension deduplication against missing IDs.

Expand Down
122 changes: 122 additions & 0 deletions apps/chat/tests/test_chat_user_quota.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
# coding=utf-8
from types import SimpleNamespace
from unittest.mock import Mock, patch

from django.test import SimpleTestCase

from chat.serializers.chat import ChatSerializers
from common.exception.app_exception import AppApiException
from system_manage.models.chat_user_token_quota import ChatUserTokenQuota, QuotaType


class ChatUserQuotaAvailabilityTests(SimpleTestCase):
def test_exhausted_periodic_quota_is_rejected_before_workflow_runs(self):
"""配额已用尽时，对话在开始前就被拒绝，不会触发工作流生成内容。"""
exhausted = SimpleNamespace(
quota_type=QuotaType.PERIODIC,
token_limit=10,
used_tokens=10,
period_end=None,
check_and_reset=lambda: None,
)

data = {
"chat_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1a",
"chat_user_id": "user-1",
"chat_user_type": "CHAT_USER",
"application_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1b",
"debug": False,
"ip_address": "127.0.0.1",
"source": {"type": "ONLINE"},
}
instance = {"message": {"content": "hi"}, "stream": False}

with (
patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: exhausted)),
patch.object(ChatSerializers, "get_application", return_value=object()),
patch.object(ChatSerializers, "is_valid_chat"),
patch.object(ChatSerializers, "is_valid_intraday_access_num"),
patch.object(ChatSerializers, "chat_work_flow") as workflow,
):
serializer = ChatSerializers(data=data)
with self.assertRaises(AppApiException):
serializer.chat(instance)
workflow.assert_not_called()

def test_available_quota_lets_workflow_run(self):
"""配额未用尽时，前置校验不拦截，正常进入工作流。"""
available = SimpleNamespace(
quota_type=QuotaType.PERIODIC,
token_limit=10,
used_tokens=1,
period_end=None,
check_and_reset=lambda: None,
)

data = {
"chat_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1a",
"chat_user_id": "user-1",
"chat_user_type": "CHAT_USER",
"application_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1b",
"debug": False,
"ip_address": "127.0.0.1",
"source": {"type": "ONLINE"},
}
instance = {"message": {"content": "hi"}, "stream": False}

with (
patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: available)),
patch.object(ChatSerializers, "get_application", return_value=object()),
patch.object(ChatSerializers, "is_valid_chat"),
patch.object(ChatSerializers, "is_valid_intraday_access_num"),
patch.object(ChatSerializers, "chat_work_flow", return_value="ok") as workflow,
):
serializer = ChatSerializers(data=data)
self.assertEqual(serializer.chat(instance), "ok")
workflow.assert_called_once()

def test_check_available_over_limit_raises(self):
"""预校验：当前周期已用满则抛异常。"""
quota = SimpleNamespace(
quota_type=QuotaType.PERIODIC,
token_limit=10,
used_tokens=10,
period_end=None,
check_and_reset=lambda: None,
)

with patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: quota)):
with self.assertRaises(AppApiException):
ChatUserTokenQuota.check_available("user-1")

def test_check_available_within_limit_passes(self):
"""预校验：未用满时不拦截。"""
quota = SimpleNamespace(
quota_type=QuotaType.PERIODIC,
token_limit=10,
used_tokens=1,
period_end=None,
check_and_reset=lambda: None,
)

with patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: quota)):
ChatUserTokenQuota.check_available("user-1")

def test_recording_over_limit_does_not_raise(self):
"""记录实际用量时即使超出上限也不抛异常，并如实累计，保证对话记录能落库。"""
quota = SimpleNamespace(
quota_type=QuotaType.PERIODIC,
token_limit=10,
used_tokens=9,
total_tokens=0,
period_end=None,
check_and_reset=lambda: None,
save=Mock(),
)

with patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: quota)):
ChatUserTokenQuota.consume("user-1", 5)

self.assertEqual(quota.used_tokens, 14)
self.assertEqual(quota.total_tokens, 5)
quota.save.assert_called_once_with(update_fields=["used_tokens", "total_tokens"])
30 changes: 20 additions & 10 deletions apps/system_manage/models/chat_user_token_quota.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from common.mixins.app_model_mixin import AppModelMixin
from dateutil.relativedelta import relativedelta
from django.utils import timezone
from django.utils.translation import gettext_lazy as _


class QuotaType(models.TextChoices):
Expand Down Expand Up @@ -68,26 +69,35 @@ def check_and_reset(self):

@classmethod
def consume(cls, user_id, amount):
"""记录本次实际 Token 消耗：累计 used_tokens/total_tokens。
即使超出上限也不抛异常，配额拦截交给对话前的 check_available，保证对话记录能正常落库。
"""
if amount <= 0:
return
quota = cls.objects.filter(user_id=user_id).first()
if quota is None:
# 所有消费用户都创建统计行，匿名用户也累计使用量
quota, _ = cls.objects.get_or_create(
quota, _created = cls.objects.get_or_create(
user_id=user_id,
defaults={"quota_type": QuotaType.UNLIMITED},
)
quota.check_and_reset()
# 记录实际使用量：即使超出上限也如实累计，拦截交给下次对话前的 check_available
quota.used_tokens += amount
quota.total_tokens += amount
quota.save(update_fields=["used_tokens", "total_tokens"])

@classmethod
def check_available(cls, user_id):
"""对话前预校验：当前周期配额已用满或超限则抛 AppApiException，不写库。"""
quota = cls.objects.filter(user_id=user_id).first()
if quota is None:
return
if quota.quota_type == QuotaType.UNLIMITED:
# 不限额：只累计使用量，不校验上限
quota.used_tokens += amount
quota.total_tokens += amount
quota.save(update_fields=["used_tokens", "total_tokens"])
return
quota.check_and_reset()
if quota.used_tokens + amount > quota.token_limit:
if quota.token_limit is not None and quota.used_tokens >= quota.token_limit:
raise AppApiException(
500, _("The token quota for the current period has been exhausted. Please contact the administrator.")
500,
_("The token quota for the current period has been exhausted. Please contact the administrator."),
)
quota.used_tokens += amount
quota.total_tokens += amount
quota.save(update_fields=["used_tokens", "total_tokens"])

@shaohuzhang1 shaohuzhang1 Nov 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is an analysis of the code changes, highlighting potential issues and optimization suggestions:

1. Logic Change & Behavioral Shift

  • Old Behavior: consume used exceptions to enforce quotas during consumption. If an LLM call was being processed, it might fail mid-stream or after generation if quota was exceeded, depending on where consume is called relative to token counting.
  • New Behavior: consume now only records usage (even if over-limit). Enforcement is moved to check_available, which must be called before starting an LLM interaction.
  • Critical Requirement: This assumes that check_available is reliably called before every LLM request. If any code path misses this call, users with exhausted quotas will continue to generate responses for free (while their stats update correctly). Ensure all entry points call check_available.

2. Potential Race Conditions (Concurrency)

  • Issue: The pattern quota.used_tokens += amount followed by quota.save(...) is not atomic. In high-concurrency scenarios (multiple requests from the same user simultaneously), this can lead to lost updates (e.g., two concurrent +100 operations result in a net +100 instead of +200).
  • Suggestion: Use Django's F() expressions to perform atomic database increments: 
    from django.db.models import F

# Inside consume():
quota.check_and_reset()
quota_model = cls.objects.filter(id=quota.id).first() # Re-fetch to ensure latest state after reset
if quota_model:
    quota_model.used_tokens += amount
    quota_model.total_tokens += amount
    quota_model.save(update_fields=["used_tokens", "total_tokens"])

# Better yet, use atomic update:
cls.objects.filter(id=quota.id, user_id=user_id).update(
    used_tokens=F('used_tokens') + amount,
    total_tokens=F('total_tokens') + amount
)
    Note: Using F() requires careful handling if check_and_reset might have changed the quota_type or reset counters in between reads and writes.

3. Redundant Database Queries

  • Issue: In consume, you fetch the quota, potentially create it, call check_and_reset, then modify and save. In check_available, you fetch again. If both are called sequentially for the same user/request, you hit the DB multiple times unnecessarily.
  • Suggestion: Consider combining logic or caching the quota object within the request context if feasible. At minimum, ensure get_or_create isn't too frequent (it should be rare after initial setup).

4. Error Handling in check_available

  • Observation: check_available raises an exception only if token_limit is not None AND used_tokens >= limit.
  • Potential Issue: If token_limit is None for non-UNLIMITED types (which should ideally not happen, but might due to data inconsistency), the check passes silently. Ensure data integrity constraints prevent token_limit=None for non-unlimited quotas.

5. Code Style & Readability

  • Good: Renaming _ to _created in get_or_create improves clarity.
  • Good: Clear docstrings explaining the separation of concerns (check_available vs consume).
  • Suggestion: Add type hints for better maintainability: 
    def consume(cls, user_id: int, amount: int) -> None:
    ...
@classmethod
def check_available(cls, user_id: int) -> None:
    ...

6. Edge Case: Zero/Limit Boundary

  • Check: if quota.used_tokens >= quota.token_limit: — This blocks usage exactly at the limit. Confirm this matches business rules (i.e., is the last allowed token still usable? Old code used > limit + amount ? No, old was +amount > limit`, meaning if you had 90/100 left and requested 20, it blocked. New code blocks if you are at 100/100. If you’re at 95/100 and send a 3-token request, new code allows it (resulting in 98/100 or 98>100? Wait, new code checks before sending. So at 95, it passes. Then after response, used becomes 98. Next time passes. Eventually hits 100, then blocks. This seems reasonable.)

Summary of Recommendations:

  1. Add F() expressions for thread-safe atomic increments in consume.
  2. Ensure check_available is called in all possible entry points for LLM interactions.
  3. Consider merging check_available logic into a single method that returns status + updated quota object to reduce DB queries per request.
  4. Validate that no code path bypasses check_available.

The overall refactoring improves clarity and aligns with best practices (fail fast before expensive operations), but concurrency handling needs attention.

Loading