diff --git a/api/controllers/console/explore/completion.py b/api/controllers/console/explore/completion.py index 2f27827a0a6..6b546d183d4 100644 --- a/api/controllers/console/explore/completion.py +++ b/api/controllers/console/explore/completion.py @@ -36,6 +36,7 @@ from models import Account from models.model import AppMode, InstalledApp from services.app_generate_service import AppGenerateService from services.app_task_service import AppTaskService +from services.conversation_service import ConversationService from services.errors.llm import InvokeRateLimitError from .. import console_ns @@ -188,6 +189,15 @@ class ChatApi(InstalledAppResource): db.session.commit() try: + # Eagerly validate conversation to avoid hanging on invalid conversation_id + if payload.conversation_id: + ConversationService.get_conversation( + app_model=app_model, + conversation_id=payload.conversation_id, + user=current_user, + session=db.session(), + ) + response = AppGenerateService.generate( session=session, app_model=app_model, diff --git a/api/controllers/service_api/app/completion.py b/api/controllers/service_api/app/completion.py index 165262d8b0a..7bbe1ec5ffc 100644 --- a/api/controllers/service_api/app/completion.py +++ b/api/controllers/service_api/app/completion.py @@ -40,12 +40,14 @@ from core.errors.error import ( QuotaExceededError, ) from core.helper.trace_id_helper import get_external_trace_id, get_trace_session_id, omit_trace_session_id_from_payload +from extensions.ext_database import db from graphon.model_runtime.errors.invoke import InvokeError from libs import helper from libs.helper import UUIDStrOrEmpty from models.model import App, AppMode, EndUser from services.app_generate_service import AppGenerateService from services.app_task_service import AppTaskService +from services.conversation_service import ConversationService from services.errors.app import IsDraftWorkflowError, WorkflowIdFormatError, WorkflowNotFoundError from services.errors.llm import InvokeRateLimitError @@ -377,6 +379,15 @@ class ChatApi(Resource): streaming = _resolve_agent_app_streaming(app_mode=app_mode, response_mode=payload.response_mode) try: + # Eagerly validate conversation to avoid hanging on invalid conversation_id + if payload.conversation_id: + ConversationService.get_conversation( + app_model=app_model, + conversation_id=payload.conversation_id, + user=end_user, + session=db.session(), + ) + response = AppGenerateService.generate( session=session, app_model=app_model, diff --git a/api/tests/unit_tests/controllers/console/explore/test_completion.py b/api/tests/unit_tests/controllers/console/explore/test_completion.py index 53d9badc5b8..64c87043889 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_completion.py +++ b/api/tests/unit_tests/controllers/console/explore/test_completion.py @@ -1,3 +1,4 @@ +import uuid from inspect import unwrap from unittest.mock import MagicMock, PropertyMock, patch @@ -310,6 +311,37 @@ class TestChatApi: with pytest.raises(completion_module.NotFound): method(api, MagicMock(), user, chat_app) + def test_invalid_conversation_id_fails_fast_as_not_found(self, app: Flask, chat_app, user) -> None: + # A nonexistent conversation_id must fail fast as 404, before the streaming + # generator is created. Previously the lookup only ran inside the generator, + # so an invalid id surfaced as a hang instead of a clean error. + payload_patch = patch.object( + type(completion_module.console_ns), + "payload", + new_callable=PropertyMock, + return_value={"inputs": {}, "query": "hi", "conversation_id": str(uuid.uuid4())}, + ) + generate_mock = MagicMock(return_value={"ok": True}) + + api = completion_module.ChatApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={}), + payload_patch, + patch.object( + completion_module.ConversationService, + "get_conversation", + side_effect=completion_module.services.errors.conversation.ConversationNotExistsError(), + ), + patch.object(completion_module.AppGenerateService, "generate", generate_mock), + ): + with pytest.raises(completion_module.NotFound): + method(api, MagicMock(), user, chat_app) + + # The lookup must run before generation, so the generator is never started. + generate_mock.assert_not_called() + def test_app_unavailable_chat(self, app: Flask, chat_app, user, payload_patch): api = completion_module.ChatApi() method = unwrap(api.post) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_completion.py b/api/tests/unit_tests/controllers/service_api/app/test_completion.py index 65652594294..b9a1b364ea5 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_completion.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_completion.py @@ -42,6 +42,7 @@ from graphon.model_runtime.errors.invoke import InvokeError from models.model import App, AppMode, EndUser from services.app_generate_service import AppGenerateService from services.app_task_service import AppTaskService +from services.conversation_service import ConversationService from services.errors.app import IsDraftWorkflowError, WorkflowIdFormatError, WorkflowNotFoundError from services.errors.conversation import ConversationNotExistsError from services.errors.llm import InvokeRateLimitError @@ -554,6 +555,35 @@ class TestChatApiController: with pytest.raises(AgentNotPublishedError): handler(api, session=Mock(), app_model=app_model, end_user=end_user) + def test_invalid_conversation_id_fails_fast_as_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + # A well-formed but nonexistent conversation_id must fail fast as 404, before the + # streaming generator is created. Previously the lookup only ran inside the generator, + # so an invalid id surfaced as a hang instead of a clean error. + monkeypatch.setattr( + ConversationService, + "get_conversation", + lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), + ) + + generate_mock = Mock(return_value={"text": "unused"}) + monkeypatch.setattr(AppGenerateService, "generate", generate_mock) + + api = ChatApi() + handler = unwrap(api.post) + app_model = SimpleNamespace(mode=AppMode.CHAT.value, id="app-1") + end_user = SimpleNamespace() + + with app.test_request_context( + "/chat-messages", + method="POST", + json={"inputs": {}, "query": "hi", "conversation_id": str(uuid.uuid4())}, + ): + with pytest.raises(NotFound): + handler(api, session=Mock(), app_model=app_model, end_user=end_user) + + # The lookup must run before generation, so the generator is never started. + generate_mock.assert_not_called() + class TestChatStopApiController: def test_wrong_mode(self, app: Flask) -> None: