diff --git a/backend/apps/agent_app.py b/backend/apps/agent_app.py index 6aad1ce9e..c5947a23f 100644 --- a/backend/apps/agent_app.py +++ b/backend/apps/agent_app.py @@ -127,6 +127,8 @@ async def agent_run_api( ) except ForbiddenError as e: raise HTTPException(status_code=HTTPStatus.FORBIDDEN, detail=str(e)) from e + except AppException: + raise except ValidationError as e: raise HTTPException( status_code=HTTPStatus.UNPROCESSABLE_ENTITY, diff --git a/backend/apps/conversation_management_app.py b/backend/apps/conversation_management_app.py index 0688788d1..69d8d691e 100644 --- a/backend/apps/conversation_management_app.py +++ b/backend/apps/conversation_management_app.py @@ -13,7 +13,11 @@ OpinionRequest, RenameRequest, ) -from consts.exceptions import ConversationNotFoundError, ValidationError +from consts.exceptions import ( + AppException, + ConversationNotFoundError, + ValidationError, +) from services.conversation_management_service import ( create_new_conversation, delete_conversation_service, @@ -51,6 +55,8 @@ async def create_new_conversation_endpoint(request: ConversationRequest, authori user_id, tenant_id = get_current_user_id(authorization) conversation_data = create_new_conversation(request.title, user_id) return ConversationResponse(code=0, message="success", data=conversation_data) + except AppException: + raise except Exception as e: logging.error(f"Failed to create conversation: {str(e)}") raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=str(e)) diff --git a/backend/consts/const.py b/backend/consts/const.py index 577dfabbb..5b7723389 100644 --- a/backend/consts/const.py +++ b/backend/consts/const.py @@ -183,6 +183,8 @@ class VectorDatabaseType(str, Enum): MAX_GROUPS_PER_TENANT = 1_000 MAX_SUPER_ADMIN_COUNT = 1 MAX_ADMINS_PER_TENANT = 1_000 +MAX_CONVERSATION_TURNS = 100 +MAX_CONVERSATIONS_PER_USER = 1_000 # Invitation code type for asset administrator registration ASSET_OWNER_INVITE_CODE_TYPE = "ASSET_OWNER_INVITE" diff --git a/backend/consts/error_code.py b/backend/consts/error_code.py index 24fd9d2f5..a99053c8c 100644 --- a/backend/consts/error_code.py +++ b/backend/consts/error_code.py @@ -278,6 +278,8 @@ class ErrorCode(Enum): ErrorCode.COMMON_MISSING_REQUIRED_FIELD: 400, # Common - Rate Limit ErrorCode.COMMON_RATE_LIMIT_EXCEEDED: 429, + # Tenant resources + ErrorCode.TENANT_RESOURCE_EXCEEDED: 429, # Common - Resource ErrorCode.COMMON_RESOURCE_NOT_FOUND: 404, ErrorCode.COMMON_RESOURCE_ALREADY_EXISTS: 409, diff --git a/backend/database/conversation_db.py b/backend/database/conversation_db.py index 0a8b91c20..08c56c4a6 100644 --- a/backend/database/conversation_db.py +++ b/backend/database/conversation_db.py @@ -2,8 +2,15 @@ from datetime import datetime from typing import Any, Dict, List, Optional, TypedDict -from sqlalchemy import asc, desc, func, insert, select, update +from sqlalchemy import asc, desc, func, insert, select, text, update +from consts.const import ( + MAX_CONVERSATION_TURNS, + MAX_CONVERSATIONS_PER_USER, + MESSAGE_ROLE, +) +from consts.error_code import ErrorCode +from consts.exceptions import AppException from .client import as_dict, db_client, get_db_session from .db_models import ( ConversationMessage, @@ -102,6 +109,69 @@ def _get_effective_tenant_id(user_tenant: Dict[str, Any]) -> str: return DEFAULT_TENANT_ID +def count_user_turns(conversation_id: int) -> int: + """Count active user messages, treating each one as one conversation turn.""" + with get_db_session() as session: + statement = select(func.count(ConversationMessage.message_id)).where( + ConversationMessage.conversation_id == int(conversation_id), + ConversationMessage.message_role == MESSAGE_ROLE["USER"], + ConversationMessage.delete_flag == "N", + ) + return int(session.execute(statement).scalar_one()) + + +def _lock_limit_scope(session, scope: str) -> None: + """Serialize quota checks for one user or conversation in PostgreSQL.""" + session.execute( + text("SELECT pg_advisory_xact_lock(hashtext(:lock_key))"), + {"lock_key": scope}, + ) + + +def _enforce_user_conversation_limit(session, user_id: Optional[str]) -> None: + if not user_id: + return + + _lock_limit_scope(session, f"conversation-user-limit:{user_id}") + statement = select(func.count(ConversationRecord.conversation_id)).where( + ConversationRecord.created_by == user_id, + ConversationRecord.delete_flag == "N", + ) + conversation_count = int(session.execute(statement).scalar_one()) + if conversation_count >= MAX_CONVERSATIONS_PER_USER: + raise AppException( + ErrorCode.TENANT_RESOURCE_EXCEEDED, + "Conversation history limit reached: " + f"maximum {MAX_CONVERSATIONS_PER_USER} conversations per user", + details={"resource": "conversations", "limit": MAX_CONVERSATIONS_PER_USER}, + ) + + +def _enforce_conversation_turn_limit( + session, + conversation_id: int, + message_role: str, + user_id: Optional[str], +) -> None: + if not user_id or message_role != MESSAGE_ROLE["USER"]: + return + + _lock_limit_scope(session, f"conversation-turn-limit:{conversation_id}") + statement = select(func.count(ConversationMessage.message_id)).where( + ConversationMessage.conversation_id == conversation_id, + ConversationMessage.message_role == MESSAGE_ROLE["USER"], + ConversationMessage.delete_flag == "N", + ) + turn_count = int(session.execute(statement).scalar_one()) + if turn_count >= MAX_CONVERSATION_TURNS: + raise AppException( + ErrorCode.TENANT_RESOURCE_EXCEEDED, + "Conversation turn limit reached: " + f"maximum {MAX_CONVERSATION_TURNS} turns per conversation", + details={"resource": "conversation_turns", "limit": MAX_CONVERSATION_TURNS}, + ) + + def create_conversation(conversation_title: str, user_id: Optional[str] = None, agent_id: Optional[int] = None, chat_mode: Optional[str] = None, @@ -120,6 +190,8 @@ def create_conversation(conversation_title: str, user_id: Optional[str] = None, Dict[str, Any]: Dictionary containing complete information of the newly created conversation """ with get_db_session() as session: + _enforce_user_conversation_limit(session, user_id) + # Prepare data dictionary data = {"conversation_title": conversation_title, "delete_flag": 'N'} if agent_id is not None: @@ -180,6 +252,13 @@ def create_conversation_message(message_data: Dict[str, Any], user_id: Optional[ # Ensure conversation_id is integer type conversation_id = int(message_data['conversation_id']) message_idx = int(message_data['message_idx']) + message_role = message_data['role'] + _enforce_conversation_turn_limit( + session=session, + conversation_id=conversation_id, + message_role=message_role, + user_id=user_id, + ) minio_files = message_data.get('minio_files') # Convert minio_files to JSON string for storage @@ -189,7 +268,7 @@ def create_conversation_message(message_data: Dict[str, Any], user_id: Optional[ minio_files = json.dumps(minio_files) # Prepare data dictionary - data = {"conversation_id": conversation_id, "message_index": message_idx, "message_role": message_data['role'], + data = {"conversation_id": conversation_id, "message_index": message_idx, "message_role": message_role, "message_content": message_data['content'], "minio_files": minio_files, "opinion_flag": None, "delete_flag": 'N', "status": status} if user_id: diff --git a/backend/services/agent_service.py b/backend/services/agent_service.py index 41e2e446e..bedf27c38 100644 --- a/backend/services/agent_service.py +++ b/backend/services/agent_service.py @@ -94,6 +94,7 @@ from utils.str_utils import convert_list_to_string, convert_string_to_list from services.conversation_management_service import ( create_new_conversation, + ensure_conversation_turn_capacity, generate_conversation_title_service, # noqa: F401 - compatibility patch point get_conversation_service, get_current_run_user_message_id, @@ -3320,6 +3321,18 @@ async def run_agent_stream( resolved_user_id, ) + if ( + not agent_request.is_debug + and not resume + and not skip_user_save + and not is_new_conversation + and agent_request.conversation_id is not None + ): + ensure_conversation_turn_capacity( + conversation_id=agent_request.conversation_id, + user_id=resolved_user_id, + ) + if ( not agent_request.is_debug and not is_new_conversation diff --git a/backend/services/conversation_management_service.py b/backend/services/conversation_management_service.py index 1e6d2da56..37a44633b 100644 --- a/backend/services/conversation_management_service.py +++ b/backend/services/conversation_management_service.py @@ -6,11 +6,20 @@ from jinja2 import StrictUndefined, Template -from consts.const import LANGUAGE, MODEL_CONFIG_MAPPING, MESSAGE_ROLE, DEFAULT_EN_TITLE, DEFAULT_ZH_TITLE +from consts.const import ( + DEFAULT_EN_TITLE, + DEFAULT_ZH_TITLE, + LANGUAGE, + MAX_CONVERSATION_TURNS, + MESSAGE_ROLE, + MODEL_CONFIG_MAPPING, +) +from consts.error_code import ErrorCode from consts.model import AgentRequest, MessageRequest, MessageUnit -from consts.exceptions import ConversationNotFoundError, ValidationError +from consts.exceptions import AppException, ConversationNotFoundError, ValidationError from database.conversation_db import ( CHAT_MODE_VALUES, + count_user_turns, create_conversation, create_conversation_message, create_message_unit, @@ -106,6 +115,21 @@ def save_message(request: MessageRequest, user_id: str, tenant_id: str, return create_conversation_message(message_data_copy, user_id, status=status) +def ensure_conversation_turn_capacity(conversation_id: int, user_id: str) -> None: + """Reject a new turn before streaming starts when the conversation is full.""" + if not user_id: + return + + turn_count = count_user_turns(conversation_id) + if turn_count >= MAX_CONVERSATION_TURNS: + raise AppException( + ErrorCode.TENANT_RESOURCE_EXCEEDED, + "Conversation turn limit reached: " + f"maximum {MAX_CONVERSATION_TURNS} turns per conversation", + details={"resource": "conversation_turns", "limit": MAX_CONVERSATION_TURNS}, + ) + + def save_message_unit(message_id: int, conversation_id: int, unit_index: int, unit_type: str, unit_content: Any, user_id: Optional[str] = None, @@ -378,6 +402,8 @@ def create_new_conversation( create_kwargs["knowledge_scope"] = knowledge_scope conversation_data = create_conversation(title, user_id, **create_kwargs) return conversation_data + except AppException: + raise except Exception as e: logging.error(f"Failed to create conversation: {str(e)}") raise Exception(str(e)) diff --git a/frontend/app/[locale]/chat/internal/chatInterface.tsx b/frontend/app/[locale]/chat/internal/chatInterface.tsx index 544e4e9db..203607c54 100644 --- a/frontend/app/[locale]/chat/internal/chatInterface.tsx +++ b/frontend/app/[locale]/chat/internal/chatInterface.tsx @@ -9,6 +9,7 @@ import { useTranslation } from "react-i18next"; import { ROLE_ASSISTANT } from "@/const/agentConfig"; import { MESSAGE_ROLES } from "@/const/chatConfig"; +import { getConversationResourceLimitMessage } from "@/const/errorMessageI18n"; import { useConfig } from "@/hooks/useConfig"; import { useModelList } from "@/hooks/model/useModelList"; import { useAuthorizationContext } from "@/components/providers/AuthorizationProvider"; @@ -1034,7 +1035,9 @@ export function ChatInterface() { }); } else { log.error(t("chatInterface.errorLabel"), error); - const errorMessage = t("chatInterface.errorProcessingRequest"); + const errorMessage = + getConversationResourceLimitMessage(error, t) || + t("chatInterface.errorProcessingRequest"); setSessionMessages((prev) => { const newMessages = { ...prev }; const lastMsg = diff --git a/frontend/app/[locale]/newchat/assistant-ui/thread.tsx b/frontend/app/[locale]/newchat/assistant-ui/thread.tsx index 3c4ada572..b699f9d19 100644 --- a/frontend/app/[locale]/newchat/assistant-ui/thread.tsx +++ b/frontend/app/[locale]/newchat/assistant-ui/thread.tsx @@ -63,6 +63,7 @@ import { message } from "antd"; import type { Agent, PublishedAgent } from "@/types/agentConfig"; import { getAgentIcon } from "@/lib/chat/agentIconUtils"; import type { ModelOption } from "../ui/model-selector"; +import { getConversationResourceLimitMessage } from "@/const/errorMessageI18n"; import AutomationProposalMessage from "@/features/agentAutomation/components/AutomationProposalMessage"; import type { AgentAutomationProposalData } from "@/types/agentAutomation"; import { @@ -964,10 +965,21 @@ const ThreadScrollToBottom: FC = () => { }; const MessageError: FC = () => { + const { t } = useTranslation("common"); + const rawError = useAuiState((state) => { + const status = state.message.status; + return status?.type === "incomplete" && status.reason === "error" + ? status.error + : undefined; + }); + const localizedMessage = getConversationResourceLimitMessage(rawError, t); + return ( - + + {localizedMessage ?? undefined} + ); diff --git a/frontend/const/errorMessageI18n.ts b/frontend/const/errorMessageI18n.ts index 1c21257f7..dfb06f5a9 100644 --- a/frontend/const/errorMessageI18n.ts +++ b/frontend/const/errorMessageI18n.ts @@ -43,6 +43,112 @@ export const getI18nErrorMessage = ( ); }; +const CONVERSATION_RESOURCE_LIMIT_TRANSLATION_KEYS: Record = { + conversations: "chatInterface.conversationLimitExceeded", + conversation_turns: "chatInterface.turnLimitExceeded", +}; + +const CONVERSATION_RESOURCE_LIMIT_PATTERNS: Array<{ + resource: keyof typeof CONVERSATION_RESOURCE_LIMIT_TRANSLATION_KEYS; + pattern: RegExp; +}> = [ + { + resource: "conversations", + pattern: + /Conversation history limit reached:\s*maximum\s+(\d+)\s+conversations?\s+per user/i, + }, + { + resource: "conversation_turns", + pattern: + /Conversation turn limit reached:\s*maximum\s+(\d+)\s+turns?\s+per conversation/i, + }, +]; + +type ConversationResourceLimitData = { + resource?: unknown; + limit?: unknown; +}; + +function getConversationErrorFields(error: unknown): { + code?: unknown; + message?: string; + data?: ConversationResourceLimitData; +} { + if (typeof error === "string") { + return { message: error }; + } + + if (!error || typeof error !== "object") { + return {}; + } + + const candidate = error as { + code?: unknown; + message?: unknown; + data?: unknown; + details?: unknown; + }; + const data = + candidate.data && typeof candidate.data === "object" + ? (candidate.data as ConversationResourceLimitData) + : candidate.details && typeof candidate.details === "object" + ? (candidate.details as ConversationResourceLimitData) + : undefined; + + return { + code: candidate.code, + message: + typeof candidate.message === "string" ? candidate.message : undefined, + data, + }; +} + +/** + * Translate conversation resource-limit errors for both chat implementations. + * + * The legacy chat receives the original ApiError, while assistant-ui may + * normalize it to an Error-like object and retain only the backend message. + * Handle both shapes so the user-facing language does not depend on the chat + * implementation that initiated the request. + */ +export const getConversationResourceLimitMessage = ( + error: unknown, + t: (key: string, options?: Record) => string +): string | null => { + const { + code, + message: errorMessage, + data, + } = getConversationErrorFields(error); + if (String(code) === ErrorCode.TENANT_RESOURCE_EXCEEDED) { + const resource = typeof data?.resource === "string" ? data.resource : ""; + const limit = data?.limit; + const translationKey = + CONVERSATION_RESOURCE_LIMIT_TRANSLATION_KEYS[resource]; + if ( + translationKey && + (typeof limit === "number" || typeof limit === "string") + ) { + return t(translationKey, { limit }); + } + } + + if (!errorMessage) { + return null; + } + + for (const { resource, pattern } of CONVERSATION_RESOURCE_LIMIT_PATTERNS) { + const match = errorMessage.match(pattern); + if (match) { + return t(CONVERSATION_RESOURCE_LIMIT_TRANSLATION_KEYS[resource], { + limit: match[1], + }); + } + } + + return null; +}; + /** * Hook to get error message with i18n support. * diff --git a/frontend/public/locales/en/common.json b/frontend/public/locales/en/common.json index fba37a444..8274eb34c 100644 --- a/frontend/public/locales/en/common.json +++ b/frontend/public/locales/en/common.json @@ -59,6 +59,8 @@ "chatInterface.fileProcessingStopped": "File parsing stopped", "chatInterface.conversationStopped": "Conversation stopped", "chatInterface.errorProcessingRequest": "An error occurred while processing the request", + "chatInterface.conversationLimitExceeded": "Conversation history has reached the limit of {{limit}} per user", + "chatInterface.turnLimitExceeded": "This conversation has reached the limit of {{limit}} turns", "chatInterface.errorFetchingConversationList": "Failed to fetch conversation list during initialization:", "chatInterface.errorFetchingConversationDetailsError": "Error fetching conversation details:", "chatInterface.errorFetchingAttachmentUrl": "Failed to fetch attachment URL: {{object_name}}", diff --git a/frontend/public/locales/zh/common.json b/frontend/public/locales/zh/common.json index 9f964a542..bf7e58b4c 100644 --- a/frontend/public/locales/zh/common.json +++ b/frontend/public/locales/zh/common.json @@ -59,6 +59,8 @@ "chatInterface.fileProcessingStopped": "文件解析已停止", "chatInterface.conversationStopped": "对话已停止", "chatInterface.errorProcessingRequest": "处理请求时发生错误", + "chatInterface.conversationLimitExceeded": "历史会话数量已达到上限：每位用户最多 {{limit}} 个", + "chatInterface.turnLimitExceeded": "当前会话对话轮数已达到上限：每个会话最多 {{limit}} 轮", "chatInterface.errorFetchingConversationList": "初始化时获取对话列表失败:", "chatInterface.errorFetchingConversationDetailsError": "获取对话详情错误:", "chatInterface.errorFetchingAttachmentUrl": "获取附件URL失败: {{object_name}}", diff --git a/frontend/services/api.ts b/frontend/services/api.ts index 5b2af4ea9..7b23cfe2f 100644 --- a/frontend/services/api.ts +++ b/frontend/services/api.ts @@ -735,7 +735,8 @@ export const API_ENDPOINTS = { export class ApiError extends Error { constructor( public code: string | number, - message: string + message: string, + public data?: Record ) { super(message); this.name = "ApiError"; @@ -755,23 +756,40 @@ export const fetchWithErrorHandling = async ( // Try to parse JSON response for business error code first let errorCode = response.status; let errorMessage = `Request failed: ${response.status}`; + let errorData: Record | undefined; const errorText = await response.text(); try { - const errorData = JSON.parse(errorText); - const errorDetail = - errorData?.detail && typeof errorData.detail === "object" - ? errorData.detail - : errorData?.message && typeof errorData.message === "object" - ? errorData.message - : errorData; + const parsedError = JSON.parse(errorText); + let errorDetail = parsedError; + if (parsedError?.detail && typeof parsedError.detail === "object") { + errorDetail = parsedError.detail; + } else if ( + parsedError?.message && + typeof parsedError.message === "object" + ) { + errorDetail = parsedError.message; + } if (errorDetail?.code) { errorCode = errorDetail.code; errorMessage = errorDetail.message || errorMessage; - } else if (typeof errorData?.detail === "string") { - errorMessage = errorData.detail; - } else if (typeof errorData?.message === "string") { - errorMessage = errorData.message; + if ( + errorDetail.data && + typeof errorDetail.data === "object" && + !Array.isArray(errorDetail.data) + ) { + errorData = errorDetail.data as Record; + } else if ( + errorDetail.details && + typeof errorDetail.details === "object" && + !Array.isArray(errorDetail.details) + ) { + errorData = errorDetail.details as Record; + } + } else if (typeof parsedError?.detail === "string") { + errorMessage = parsedError.detail; + } else if (typeof parsedError?.message === "string") { + errorMessage = parsedError.message; } else { errorMessage = errorText || errorMessage; } @@ -788,13 +806,13 @@ export const fetchWithErrorHandling = async ( errorCodeStr === ErrorCode.TOKEN_INVALID ) { handleSessionExpired(); - throw new ApiError(errorCode, errorMessage); + throw new ApiError(errorCode, errorMessage, errorData); } // Handle HTTP 401 - trigger session expired modal for all unauthorized errors if (response.status === 401) { handleSessionExpired(); - throw new ApiError(errorCode, errorMessage); + throw new ApiError(errorCode, errorMessage, errorData); } // Handle custom 499 error code (client closed connection) @@ -828,7 +846,7 @@ export const fetchWithErrorHandling = async ( ); } - throw new ApiError(errorCode, errorMessage); + throw new ApiError(errorCode, errorMessage, errorData); } return response; diff --git a/test/backend/app/test_agent_app.py b/test/backend/app/test_agent_app.py index 4ca23f962..232f24066 100644 --- a/test/backend/app/test_agent_app.py +++ b/test/backend/app/test_agent_app.py @@ -435,6 +435,38 @@ async def test_agent_run_api_maps_scope_validation_to_422(mocker): assert exc_info.value.detail == "Selected knowledge bases are incompatible" +@pytest.mark.asyncio +async def test_agent_run_api_preserves_conversation_resource_limit(mocker): + from consts.error_code import ErrorCode + from consts.exceptions import AppException + from consts.model import AgentRequest + from starlette.requests import Request + + from apps.agent_app import agent_run_api + + mocker.patch( + "apps.agent_app.run_agent_stream", + new_callable=AsyncMock, + side_effect=AppException( + ErrorCode.TENANT_RESOURCE_EXCEEDED, + "Conversation turn limit reached", + details={"resource": "conversation_turns", "limit": 100}, + ), + ) + request = AgentRequest(agent_id=1, conversation_id=123, query="test query", history=[]) + + with pytest.raises(AppException) as exc_info: + await agent_run_api( + agent_request=request, + http_request=Request({"type": "http", "headers": []}), + authorization="Bearer token", + resume=False, + ) + + assert exc_info.value.http_status == 429 + assert exc_info.value.details == {"resource": "conversation_turns", "limit": 100} + + def test_agent_stop_api_success(mocker, mock_conversation_id): """Test agent_stop_api success case.""" mock_get_user_id = mocker.patch("apps.agent_app.get_current_user_id") diff --git a/test/backend/app/test_conversation_management_app.py b/test/backend/app/test_conversation_management_app.py index 71e7867fb..bc11ca401 100644 --- a/test/backend/app/test_conversation_management_app.py +++ b/test/backend/app/test_conversation_management_app.py @@ -49,7 +49,8 @@ get_message_id_endpoint, update_conversation_knowledge_scope_endpoint, ) -from consts.exceptions import ValidationError +from consts.error_code import ErrorCode +from consts.exceptions import AppException, ValidationError # ----------------------------- @@ -151,6 +152,25 @@ async def test_create_new_conversation_failure(conversation_mocks): conversation_mocks['logging'].error.assert_called_once() +@pytest.mark.asyncio +async def test_create_new_conversation_resource_limit_preserves_structured_error(conversation_mocks): + conversation_mocks['get_current_user_id'].return_value = ("user_id", "tenant_id") + conversation_mocks['create_new_convo'].side_effect = AppException( + ErrorCode.TENANT_RESOURCE_EXCEEDED, + "Conversation history limit reached", + details={"resource": "conversations", "limit": 1_000}, + ) + + request_obj = MagicMock() + request_obj.title = "New Conversation" + + with pytest.raises(AppException) as exc_info: + await create_new_conversation_endpoint(request_obj, authorization="Bearer token") + + assert exc_info.value.http_status == 429 + assert exc_info.value.details == {"resource": "conversations", "limit": 1_000} + + @pytest.mark.asyncio async def test_list_conversations_success(conversation_mocks): """Verify successful retrieval of conversation list""" diff --git a/test/backend/database/test_conversation_db.py b/test/backend/database/test_conversation_db.py index c36f1c900..13278eb73 100644 --- a/test/backend/database/test_conversation_db.py +++ b/test/backend/database/test_conversation_db.py @@ -27,6 +27,7 @@ def _reset_captured(): sa_mod.desc = MagicMock(name="desc") sa_mod.func = MagicMock(name="func") sa_mod.select = MagicMock(name="select") +sa_mod.text = MagicMock(name="text") def _create_insert_mock(): @@ -226,6 +227,9 @@ def _add_update_tracking(data, user_id): update_message_unit_status, update_conversation_knowledge_scope, ) +import backend.database.conversation_db as conversation_db +from consts.error_code import ErrorCode +from consts.exceptions import AppException @pytest.fixture(autouse=True) @@ -506,7 +510,7 @@ def test_create_conversation_success(monkeypatch, mock_session_ctx): assert result["agent_id"] == 7 assert result["create_time"] == 1234567890 assert result["update_time"] == 1234567890 - session.execute.assert_called_once() + assert session.execute.call_count == 3 assert _captured_insert_values["agent_id"] == 7 @@ -529,6 +533,54 @@ def test_create_conversation_without_user_id(monkeypatch, mock_session_ctx): session.execute.assert_called_once() +def test_user_conversation_limit_allows_value_below_limit(): + session = MagicMock() + count_result = MagicMock() + count_result.scalar_one.return_value = 999 + session.execute.side_effect = [MagicMock(), count_result] + + conversation_db._enforce_user_conversation_limit(session, "user-1") + + +def test_user_conversation_limit_rejects_value_at_limit(): + session = MagicMock() + count_result = MagicMock() + count_result.scalar_one.return_value = 1_000 + session.execute.side_effect = [MagicMock(), count_result] + + with pytest.raises(AppException) as exc_info: + conversation_db._enforce_user_conversation_limit(session, "user-1") + + assert exc_info.value.error_code == ErrorCode.TENANT_RESOURCE_EXCEEDED + assert exc_info.value.details == {"resource": "conversations", "limit": 1_000} + + +def test_conversation_turn_limit_allows_value_below_limit(): + session = MagicMock() + count_result = MagicMock() + count_result.scalar_one.return_value = 99 + session.execute.side_effect = [MagicMock(), count_result] + + conversation_db._enforce_conversation_turn_limit( + session, conversation_id=1, message_role="user", user_id="user-1" + ) + + +def test_conversation_turn_limit_rejects_value_at_limit(): + session = MagicMock() + count_result = MagicMock() + count_result.scalar_one.return_value = 100 + session.execute.side_effect = [MagicMock(), count_result] + + with pytest.raises(AppException) as exc_info: + conversation_db._enforce_conversation_turn_limit( + session, conversation_id=1, message_role="user", user_id="user-1" + ) + + assert exc_info.value.error_code == ErrorCode.TENANT_RESOURCE_EXCEEDED + assert exc_info.value.details == {"resource": "conversation_turns", "limit": 100} + + # ============================================================================= # Tests for create_conversation_message # ============================================================================= diff --git a/test/backend/services/test_conversation_management_service.py b/test/backend/services/test_conversation_management_service.py index dcddaf940..c0027d054 100644 --- a/test/backend/services/test_conversation_management_service.py +++ b/test/backend/services/test_conversation_management_service.py @@ -206,6 +206,7 @@ def validate(self): pass call_llm_for_title, update_conversation_title, create_new_conversation, + ensure_conversation_turn_capacity, get_conversation_service, get_conversation_list_service, rename_conversation_service, @@ -218,6 +219,8 @@ def validate(self): pass update_message_opinion_service, get_message_id_by_index_impl ) +from consts.error_code import ErrorCode +from consts.exceptions import AppException class TestConversationManagementService(unittest.TestCase): @@ -521,6 +524,30 @@ def test_create_new_conversation(self, mock_create_conversation): mock_create_conversation.assert_called_once_with( "New Chat", self.user_id, agent_id=None, chat_mode=None) + @patch('backend.services.conversation_management_service.count_user_turns', return_value=99) + def test_ensure_conversation_turn_capacity_allows_below_limit(self, mock_count_user_turns): + ensure_conversation_turn_capacity(123, self.user_id) + + mock_count_user_turns.assert_called_once_with(123) + + @patch('backend.services.conversation_management_service.count_user_turns') + def test_ensure_conversation_turn_capacity_skips_missing_user(self, mock_count_user_turns): + ensure_conversation_turn_capacity(123, "") + + mock_count_user_turns.assert_not_called() + + @patch('backend.services.conversation_management_service.count_user_turns', return_value=100) + def test_ensure_conversation_turn_capacity_rejects_at_limit(self, mock_count_user_turns): + with self.assertRaises(AppException) as exc_info: + ensure_conversation_turn_capacity(123, self.user_id) + + self.assertEqual(exc_info.exception.error_code, ErrorCode.TENANT_RESOURCE_EXCEEDED) + self.assertEqual( + exc_info.exception.details, + {"resource": "conversation_turns", "limit": 100}, + ) + mock_count_user_turns.assert_called_once_with(123) + @patch('backend.services.conversation_management_service.create_conversation') def test_create_new_conversation_with_agent_id(self, mock_create_conversation): # Setup