From ad4121b0d8efc3c9e8b5a88798f35e4aca418208 Mon Sep 17 00:00:00 2001 From: wwq Date: Tue, 21 Apr 2026 20:48:06 +0800 Subject: [PATCH] fix(api): fix API Key rate limiting and terminal user quota checks - Revert API Key rate limit handling to throw an error instead of auto-capping when exceeding the plan limit. - Optimize terminal user quota check logic to validate only during new user creation, avoiding redundant checks. - Add method to query terminal users by `workspace_id` and `other_id`. --- api/app/controllers/api_key_controller.py | 2 ++ .../controllers/public_share_controller.py | 26 +++++++++++++++++-- api/app/repositories/end_user_repository.py | 11 ++++++++ api/app/services/api_key_service.py | 14 +++++++--- 4 files changed, 47 insertions(+), 6 deletions(-) diff --git a/api/app/controllers/api_key_controller.py b/api/app/controllers/api_key_controller.py index dce8450d..6e414276 100644 --- a/api/app/controllers/api_key_controller.py +++ b/api/app/controllers/api_key_controller.py @@ -167,6 +167,8 @@ def update_api_key( return success(data=api_key_schema.ApiKey.model_validate(api_key), msg="API Key 更新成功") + except BusinessException: + raise except Exception as e: logger.error(f"未知错误: {str(e)}", extra={ "api_key_id": str(api_key_id), diff --git a/api/app/controllers/public_share_controller.py b/api/app/controllers/public_share_controller.py index 049535b5..486854ba 100644 --- a/api/app/controllers/public_share_controller.py +++ b/api/app/controllers/public_share_controller.py @@ -219,9 +219,20 @@ def list_conversations( end_user_repo = EndUserRepository(db) app_service = AppService(db) app = app_service._get_app_or_404(share.app_id) + workspace_id = app.workspace_id + + # 仅在新建终端用户时检查配额 + existing_end_user = end_user_repo.get_end_user_by_other_id(workspace_id=workspace_id, other_id=other_id) + if existing_end_user is None: + from app.core.quota_manager import _check_quota + from app.models.workspace_model import Workspace + ws = db.query(Workspace).filter(Workspace.id == workspace_id).first() + if ws: + _check_quota(db, ws.tenant_id, "end_user_quota", "end_user") + new_end_user = end_user_repo.get_or_create_end_user( app_id=share.app_id, - workspace_id=app.workspace_id, + workspace_id=workspace_id, other_id=other_id ) logger.debug(new_end_user.id) @@ -309,7 +320,6 @@ def get_conversation( "/chat", summary="发送消息（支持流式和非流式）" ) -@check_end_user_quota async def chat( payload: conversation_schema.ChatRequest, share_data: ShareTokenData = Depends(get_share_user_id), @@ -350,6 +360,18 @@ async def chat( app_service = AppService(db) app = app_service._get_app_or_404(share.app_id) workspace_id = app.workspace_id + + # 仅在新建终端用户时检查配额，已有用户复用不受限制 + existing_end_user = end_user_repo.get_end_user_by_other_id(workspace_id=workspace_id, other_id=other_id) + logger.info(f"终端用户配额检查: workspace_id={workspace_id}, other_id={other_id}, existing={existing_end_user is not None}") + if existing_end_user is None: + from app.core.quota_manager import _check_quota + from app.models.workspace_model import Workspace + ws = db.query(Workspace).filter(Workspace.id == workspace_id).first() + if ws: + logger.info(f"新终端用户，执行配额检查: tenant_id={ws.tenant_id}") + _check_quota(db, ws.tenant_id, "end_user_quota", "end_user") + new_end_user = end_user_repo.get_or_create_end_user( app_id=share.app_id, workspace_id=workspace_id, diff --git a/api/app/repositories/end_user_repository.py b/api/app/repositories/end_user_repository.py index aad80707..aba4034f 100644 --- a/api/app/repositories/end_user_repository.py +++ b/api/app/repositories/end_user_repository.py @@ -66,6 +66,17 @@ class EndUserRepository: db_logger.error(f"查询宿主 {end_user_id} 时出错: {str(e)}") raise + def get_end_user_by_other_id(self, workspace_id: uuid.UUID, other_id: str) -> Optional["EndUser"]: + """按 workspace_id + other_id 查找终端用户，不存在返回 None""" + return ( + self.db.query(EndUser) + .filter( + EndUser.workspace_id == workspace_id, + EndUser.other_id == other_id + ) + .first() + ) + def get_or_create_end_user( self, app_id: uuid.UUID, diff --git a/api/app/services/api_key_service.py b/api/app/services/api_key_service.py index 5143ac3e..4856365a 100644 --- a/api/app/services/api_key_service.py +++ b/api/app/services/api_key_service.py @@ -51,7 +51,7 @@ class ApiKeyService: if existing: raise BusinessException(f"API Key 名称 {data.name} 已存在", BizCode.API_KEY_DUPLICATE_NAME) - # 若 rate_limit 超过租户套餐的 api_ops_rate_limit，自动截断到套餐上限 + # 若 rate_limit 超过租户套餐的 api_ops_rate_limit，直接报错 from app.models.workspace_model import Workspace from app.core.quota_manager import get_api_ops_rate_limit @@ -59,7 +59,10 @@ class ApiKeyService: if workspace: tenant_api_ops_limit = get_api_ops_rate_limit(db, workspace.tenant_id) if tenant_api_ops_limit and data.rate_limit > tenant_api_ops_limit: - data.rate_limit = tenant_api_ops_limit + raise BusinessException( + f"API Key QPS 不能超过套餐上限 {tenant_api_ops_limit}", + BizCode.BAD_REQUEST + ) # 生成 API Key api_key = generate_api_key(data.type) @@ -162,7 +165,7 @@ class ApiKeyService: if existing: raise BusinessException(f"API Key 名称 {data.name} 已存在", BizCode.API_KEY_DUPLICATE_NAME) - # 若 rate_limit 超过租户套餐的 api_ops_rate_limit，自动截断到套餐上限 + # 若 rate_limit 超过租户套餐的 api_ops_rate_limit，直接报错 if data.rate_limit is not None: from app.models.workspace_model import Workspace from app.core.quota_manager import get_api_ops_rate_limit @@ -171,7 +174,10 @@ class ApiKeyService: if workspace: tenant_api_ops_limit = get_api_ops_rate_limit(db, workspace.tenant_id) if tenant_api_ops_limit and data.rate_limit > tenant_api_ops_limit: - data.rate_limit = tenant_api_ops_limit + raise BusinessException( + f"API Key QPS 不能超过套餐上限 {tenant_api_ops_limit}", + BizCode.BAD_REQUEST + ) update_data = data.model_dump(exclude_unset=True) ApiKeyRepository.update(db, api_key_id, update_data)