refactor(api): bound suggested question runtime sessions (#42968)

api/services/message_service.py

@@ -1,5 +1,7 @@
 import logging
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from datetime import datetime
 from typing import cast
 
 from sqlalchemy import select
@@ -11,7 +13,7 @@ from core.app.entities.app_invoke_entities import InvokeFrom, get_credit_usage_a
 from core.llm_generator.llm_generator import LLMGenerator
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_context import use_credit_usage_metadata
-from core.model_manager import ModelManager
+from core.model_manager import ModelInstance, ModelManager
 from core.ops.entities.trace_entity import TraceTaskName
 from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
 from core.ops.utils import measure_time
@@ -65,6 +67,23 @@ def attach_message_extra_contents(messages: Sequence[Message]) -> None:
         message.set_extra_contents([content.model_dump(mode="json", exclude_none=True) for content in contents])
 
 
+@dataclass(frozen=True, slots=True)
+class SuggestedQuestionsContext:
+    """Validated configuration and loaded conversation for the legacy history reader.
+
+    The conversation's scalar fields remain usable after the preparation session
+    closes. History loading must use an explicit session, never ORM properties.
+    """
+
+    app_id: str
+    tenant_id: str
+    app_mode: str
+    message_id: str
+    conversation: Conversation
+    instruction_prompt: str | None
+    model_config: object | None
+
+
 class MessageService:
     @classmethod
     def _get_agent_suggested_questions_config(
@@ -336,7 +355,7 @@ class MessageService:
         return message
 
     @classmethod
-    def get_suggested_questions_after_answer(
+    def prepare_suggested_questions_after_answer(
         cls,
         app_model: App,
         user: Account | EndUser | None,
@@ -344,7 +363,8 @@ class MessageService:
         invoke_from: InvokeFrom,
         *,
         session: Session,
-    ) -> list[str]:
+    ) -> SuggestedQuestionsContext | None:
+        """Read authorization and feature configuration; None means no published workflow."""
         if not user:
             raise ValueError("user cannot be None")
 
@@ -354,7 +374,6 @@ class MessageService:
             app_model=app_model, conversation_id=message.conversation_id, user=user, session=session
         )
 
-        model_manager = ModelManager.for_tenant(tenant_id=app_model.tenant_id)
         suggested_questions_after_answer_config: SuggestedQuestionsAfterAnswerConfig = {"enabled": False}
 
         if app_model.mode == AppMode.ADVANCED_CHAT:
@@ -365,7 +384,7 @@ class MessageService:
                 workflow = workflow_service.get_published_workflow(app_model=app_model, session=session)
 
             if workflow is None:
-                return []
+                return None
 
             app_config = AdvancedChatAppConfigManager.get_app_config(app_model=app_model, workflow=workflow)
 
@@ -411,46 +430,80 @@ class MessageService:
             if suggested_questions_after_answer_config.get("enabled", False) is False:
                 raise SuggestedQuestionsAfterAnswerDisabledError()
 
-        try:
-            model_instance = model_manager.get_default_model_instance(
-                tenant_id=app_model.tenant_id,
-                model_type=ModelType.LLM,
-            )
-        except Exception:
-            logger.exception("Failed to resolve the history model for suggested questions")
-            return []
-
-        # get memory of conversation (read-only)
-        memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance)
-
-        histories = memory.get_history_prompt_text(
-            max_token_limit=3000,
-            message_limit=3,
-        )
-
         instruction_prompt = suggested_questions_after_answer_config.get("prompt")
         if not isinstance(instruction_prompt, str) or not instruction_prompt.strip():
             instruction_prompt = None
-
-        configured_model = suggested_questions_after_answer_config.get("model")
-        with (
-            measure_time() as timer,
-            use_credit_usage_metadata({"app_type": get_credit_usage_app_type(app_model.mode)}),
-        ):
-            questions_sequence = LLMGenerator.generate_suggested_questions_after_answer(
-                tenant_id=app_model.tenant_id,
-                histories=histories,
-                instruction_prompt=instruction_prompt,
-                model_config=configured_model,
-            )
-            questions: list[str] = list(questions_sequence)
-
-        # get tracing instance
-        trace_manager = TraceQueueManager(app_id=app_model.id)
-        trace_manager.add_trace_task(
-            TraceTask(
-                TraceTaskName.SUGGESTED_QUESTION_TRACE, message_id=message_id, suggested_question=questions, timer=timer
-            )
+        return SuggestedQuestionsContext(
+            app_id=app_model.id,
+            tenant_id=app_model.tenant_id,
+            app_mode=app_model.mode,
+            message_id=message_id,
+            conversation=conversation,
+            instruction_prompt=instruction_prompt,
+            model_config=suggested_questions_after_answer_config.get("model"),
         )
 
+    @staticmethod
+    def get_suggested_questions_history_model(*, tenant_id: str) -> ModelInstance | None:
+        """Resolve the history token counter; None preserves the no-model fallback."""
+        model_manager = ModelManager.for_tenant(tenant_id=tenant_id)
+        try:
+            return model_manager.get_default_model_instance(tenant_id=tenant_id, model_type=ModelType.LLM)
+        except Exception:
+            logger.exception("Failed to resolve the history model for suggested questions")
+            return None
+
+    @classmethod
+    def get_suggested_questions_after_answer(
+        cls,
+        app_model: App,
+        user: Account | EndUser | None,
+        message_id: str,
+        invoke_from: InvokeFrom,
+        *,
+        session: Session,
+    ) -> list[str]:
+        """Compatibility entry point; the caller still owns its request session.
+
+        Admitted callers use MessageSuggestedQuestionsRuntime to release query
+        and model-resolution sessions before attachment and provider I/O.
+        """
+        context = cls.prepare_suggested_questions_after_answer(
+            app_model=app_model, user=user, message_id=message_id, invoke_from=invoke_from, session=session
+        )
+        if context is None:
+            return []
+        model_instance = cls.get_suggested_questions_history_model(tenant_id=context.tenant_id)
+        if model_instance is None:
+            return []
+        memory = TokenBufferMemory(conversation=context.conversation, model_instance=model_instance)
+        histories = memory.get_history_prompt_text(max_token_limit=3000, message_limit=3)
+
+        with (
+            measure_time() as timer,
+            use_credit_usage_metadata({"app_type": get_credit_usage_app_type(context.app_mode)}),
+        ):
+            questions = list(
+                LLMGenerator.generate_suggested_questions_after_answer(
+                    tenant_id=context.tenant_id,
+                    histories=histories,
+                    instruction_prompt=context.instruction_prompt,
+                    model_config=context.model_config,
+                )
+            )
+        cls.trace_suggested_questions(context=context, questions=questions, timer=timer)
         return questions
+
+    @staticmethod
+    def trace_suggested_questions(
+        *, context: SuggestedQuestionsContext, questions: list[str], timer: Mapping[str, datetime | None]
+    ) -> None:
+        trace_manager = TraceQueueManager(app_id=context.app_id)
+        trace_manager.add_trace_task(
+            TraceTask(
+                TraceTaskName.SUGGESTED_QUESTION_TRACE,
+                message_id=context.message_id,
+                suggested_question=questions,
+                timer=timer,
+            )
+        )

api/services/message_suggested_questions_adapters.py

@@ -1,7 +1,8 @@
-"""Bridge admitted apps and actors to the existing suggested-question runtime.
+"""Run suggested questions with bounded query and model-resolution sessions.
 
-TODO: Remove this bridge when message configuration, history and model lookups
-share a framework-neutral application boundary. Reuse their current policy here.
+Keep legacy configuration, history and provider policy with their existing
+owners. This bridge controls their lifetimes until those owners accept only
+framework-neutral inputs; it never closes the caller's scoped session.
 """
 
 from typing import override
@@ -10,7 +11,11 @@ from flask import current_app
 from sqlalchemy import select
 from sqlalchemy.orm import Session, sessionmaker
 
-from core.app.entities.app_invoke_entities import InvokeFrom
+from core.app.entities.app_invoke_entities import InvokeFrom, get_credit_usage_app_type
+from core.llm_generator.llm_generator import LLMGenerator
+from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_context import use_credit_usage_metadata
+from core.ops.utils import measure_time
 from extensions.ext_database import db
 from models import Account, App
 from models.model import EndUser
@@ -65,14 +70,60 @@ class MessageSuggestedQuestionsRuntime(MessageSuggestedQuestions):
                         f"End user {actor.end_user_id} does not exist for app {app_id} in tenant {app_owner_tenant_id}"
                     )
 
-        # Legacy history/model lookups still use db.session internally. Give them
-        # one scoped session, released on success or failure without removing the
-        # caller's session. They may still hold it during provider I/O.
-        with current_app.app_context():
-            return MessageService.get_suggested_questions_after_answer(
+            context = MessageService.prepare_suggested_questions_after_answer(
                 app_model=app_model,
                 user=user,
                 message_id=message_id,
                 invoke_from=InvokeFrom(actor.invoke_from),
-                session=db.session(),
+                session=session,
             )
+        if context is None:
+            return []
+
+        # Legacy model resolution, quota and tracing use db.session and may
+        # commit. Own one isolated scope so none can commit or close the caller's
+        # session. Release its sessions explicitly between phases; teardown also
+        # handles early returns and exceptions.
+        with current_app.app_context():
+            history_model = MessageService.get_suggested_questions_history_model(tenant_id=context.tenant_id)
+            db.session.remove()
+            if history_model is None:
+                return []
+
+            # Resolve the history model first to preserve the no-model fallback
+            # even when old history contains invalid file configuration.
+            with self._session_factory(expire_on_commit=False) as session:
+                history = TokenBufferMemory.load_history(
+                    conversation=context.conversation,
+                    app_record=app_model,
+                    session=session,
+                    message_limit=3,
+                )
+
+            histories = history.get_prompt_text(model_instance=history_model, max_token_limit=3000)
+
+            with (
+                measure_time() as timer,
+                use_credit_usage_metadata({"app_type": get_credit_usage_app_type(context.app_mode)}),
+            ):
+                model = LLMGenerator.prepare_suggested_questions_model(
+                    tenant_id=context.tenant_id, model_config=context.model_config
+                )
+                db.session.remove()
+                questions = (
+                    list(
+                        LLMGenerator.invoke_suggested_questions_after_answer(
+                            prepared_model=model,
+                            histories=histories,
+                            instruction_prompt=context.instruction_prompt,
+                        )
+                    )
+                    if model is not None
+                    else []
+                )
+
+            # Generation can return [] after catching an error. Tracing must not
+            # inherit an open or failed transaction from that invocation.
+            db.session.remove()
+            MessageService.trace_suggested_questions(context=context, questions=questions, timer=timer)
+            return questions

api/tests/unit_tests/controllers/console/explore/test_message.py

@@ -21,8 +21,9 @@ from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotIni
 from core.model_manager import ModelInstance
 from core.ops.ops_trace_manager import TraceTask
 from extensions.ext_database import db
-from graphon.model_runtime.entities import PromptMessage
-from graphon.model_runtime.entities.model_entities import ModelType
+from graphon.model_runtime.entities import AssistantPromptMessage, PromptMessage
+from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
+from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from graphon.model_runtime.errors.invoke import InvokeError
 from models import Account, App, AppMode, InstalledApp
 from models.enums import ConversationFromSource, ConversationStatus, FeedbackFromSource, FeedbackRating
@@ -198,29 +199,59 @@ class _QuestionProvider:
     messages: _Messages
     message: Message
     legacy_sessions: list[Session] = field(default_factory=list)
-    histories: list[str] = field(default_factory=list)
+    prompts: list[str] = field(default_factory=list)
+    io_sessions: list[tuple[str, bool]] = field(default_factory=list)
+    token_failure: bool = False
     traces: list[TraceTask] = field(default_factory=list)
 
     def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance:
         assert tenant_id == self.messages.harness.target_app.tenant_id
         assert model_type == ModelType.LLM
         self.messages.assert_sessions_closed()
+        assert db.session.get(App, self.messages.harness.target_app.id) is not None
         return cast(ModelInstance, self)
 
+    def record_io(self, stage: str) -> None:
+        self.io_sessions.append(
+            (
+                stage,
+                all(
+                    not session.in_transaction() and not session.identity_map
+                    for session in self.messages.sessions + self.legacy_sessions
+                ),
+            )
+        )
+
     def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int:
+        self.record_io("tokens")
+        if self.token_failure:
+            raise ValueError("Token counter failed")
         return len(prompt_messages)
 
-    def generate(
-        self, *, tenant_id: str, histories: str, instruction_prompt: str | None, model_config: object
-    ) -> list[str]:
-        assert tenant_id == self.messages.harness.target_app.tenant_id
-        assert instruction_prompt == "Published questions"
-        assert model_config is None
-        self.messages.assert_sessions_closed()
-        self.histories.append(histories)
+    def get_model_schema(self) -> AIModelEntity:
+        self.record_io("schema")
+        return AIModelEntity.model_construct(parameter_rules=[])
+
+    def invoke_llm(
+        self,
+        *,
+        prompt_messages: list[PromptMessage],
+        model_parameters: Mapping[str, object],
+        stop: list[str],
+        stream: bool,
+    ) -> LLMResult:
+        self.record_io("invoke")
+        assert stream is False
+        assert stop == []
+        assert model_parameters["temperature"] == 0.0
+        self.prompts.append(prompt_messages[0].get_text_content())
         if self.messages.external_error is not None:
             raise self.messages.external_error
-        return self.messages.questions
+        return LLMResult(
+            model="questions-model",
+            message=AssistantPromptMessage(content=json.dumps(self.messages.questions)),
+            usage=LLMUsage.empty_usage(),
+        )
 
     def add_trace_task(self, task: TraceTask) -> None:
         self.traces.append(task)
@@ -228,6 +259,7 @@ class _QuestionProvider:
     def assert_closed(self) -> None:
         self.messages.assert_sessions_closed()
         assert all(not session.in_transaction() and not session.identity_map for session in self.legacy_sessions)
+        assert all(closed for _stage, closed in self.io_sessions), self.io_sessions
 
 
 @pytest.fixture
@@ -256,7 +288,6 @@ def real_questions(
         return provider
 
     monkeypatch.setattr(message_module.ModelManager, "for_tenant", model_manager)
-    monkeypatch.setattr(message_module.LLMGenerator, "generate_suggested_questions_after_answer", provider.generate)
     monkeypatch.setattr(message_module, "TraceQueueManager", trace_manager)
     messages.services = replace(
         messages.services,
@@ -668,14 +699,13 @@ def test_shared_suggested_questions_uses_owner_tenant_and_releases_sessions(
     assert messages.harness.target_app.tenant_id != messages.harness.installed_app.tenant_id
     messages.external_error = failure
     response = messages.request("suggested-questions", message_id=real_questions.message.id)
-    if failure is None:
-        _assert_json_response(response, status=200, body={"data": messages.questions})
-        assert len(real_questions.traces) == 1
-    else:
-        _error(response, status=400, code="completion_request_error", message="Provider rejected input")
-        assert not real_questions.traces
-    assert real_questions.histories == ["Human: Hello\nAssistant: 你好"]
-    assert real_questions.legacy_sessions
+    _assert_json_response(response, status=200, body={"data": messages.questions if failure is None else []})
+    assert len(real_questions.traces) == 1
+    assert len(real_questions.prompts) == 1
+    assert real_questions.prompts[0].startswith("Human: Hello\nAssistant: 你好\nPublished questions\n")
+    assert len(real_questions.legacy_sessions) == 2
+    assert real_questions.legacy_sessions[0] is not real_questions.legacy_sessions[1]
+    assert [stage for stage, _closed in real_questions.io_sessions] == ["tokens", "schema", "invoke"]
     real_questions.assert_closed()
     messages.assert_unused()
 
@@ -694,7 +724,7 @@ def test_shared_suggested_questions_rejects_other_account_history(
         status=404,
         code="message_not_found" if entity == "message" else "conversation_not_found",
     )
-    assert not real_questions.histories
+    assert not real_questions.prompts
     real_questions.assert_closed()
 
 
@@ -727,7 +757,7 @@ def test_shared_suggested_questions_rejects_stale_admitted_app_and_actor(
         code="unauthorized" if change == "account" else "app_unavailable",
         message="Account no longer exists." if change == "account" else None,
     )
-    assert not real_questions.histories
+    assert not real_questions.prompts
     assert not real_questions.legacy_sessions
     real_questions.assert_closed()
 
@@ -753,7 +783,7 @@ def test_suggested_questions_uses_admitted_installation_until_next_request(real_
         status=404,
         code="installed_app_not_found",
     )
-    assert len(real_questions.histories) == 1
+    assert len(real_questions.prompts) == 1
     real_questions.assert_closed()
 
 
@@ -804,3 +834,17 @@ def test_generation_resource_errors_remain_specific(
     messages.external_error = failure
     _error(messages.request(operation), status=status, code=code)
     messages.assert_unused()
+
+
+def test_shared_suggested_questions_token_failure_releases_sessions(real_questions: _QuestionProvider) -> None:
+    real_questions.token_failure = True
+    _error(
+        real_questions.messages.request("suggested-questions", message_id=real_questions.message.id),
+        status=500,
+        code="internal_server_error",
+    )
+    assert [stage for stage, _closed in real_questions.io_sessions] == ["tokens"]
+    assert len(real_questions.legacy_sessions) == 1
+    assert not real_questions.prompts
+    assert not real_questions.traces
+    real_questions.assert_closed()

api/tests/unit_tests/controllers/console/explore/test_trial_suggested_questions.py

@@ -1,6 +1,7 @@
 """Trial suggested questions through HTTP, real ownership queries, and SQLite."""
 
-from collections.abc import Callable, Iterator, Sequence
+import json
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from dataclasses import dataclass, field
 from decimal import Decimal
 from typing import cast
@@ -24,8 +25,9 @@ from core.ops.ops_trace_manager import TraceTask
 from enums import DeploymentEdition
 from extensions.ext_database import db
 from extensions.ext_login import DifyLoginManager, unauthorized_handler
-from graphon.model_runtime.entities import PromptMessage
-from graphon.model_runtime.entities.model_entities import ModelType
+from graphon.model_runtime.entities import AssistantPromptMessage, PromptMessage
+from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
+from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from graphon.model_runtime.errors.invoke import InvokeError
 from libs.external_api import ExternalApi
 from models import Account, AccountTrialAppRecord, App, AppMode, Conversation, Message, Tenant, TrialApp
@@ -57,8 +59,11 @@ class _Features:
 @dataclass
 class _Provider:
     read_sessions: list[Session]
+    app_id: str
+    model_sessions: list[Session] = field(default_factory=list)
+    token_calls: int = 0
     tenant_ids: list[str] = field(default_factory=list)
-    histories: list[str] = field(default_factory=list)
+    prompts: list[str] = field(default_factory=list)
     traces: list[TraceTask] = field(default_factory=list)
     questions: list[str] = field(default_factory=lambda: ["Next question?", "Another question?"])
     failure: Exception | None = None
@@ -68,27 +73,52 @@ class _Provider:
     def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance:
         assert model_type == ModelType.LLM
         self.tenant_ids.append(tenant_id)
+        session = db.session()
+        assert session.get(App, self.app_id) is not None
+        self.model_sessions.append(session)
         if self.missing_model:
             raise ProviderTokenNotInitError("No default model")
         return cast(ModelInstance, self)
 
-    def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int:
-        return len(prompt_messages)
+    def assert_queries_closed(self) -> None:
+        assert all(
+            not session.in_transaction() and not session.identity_map
+            for session in self.read_sessions + self.model_sessions
+        )
 
-    def generate(
-        self, *, tenant_id: str, histories: str, instruction_prompt: str | None, model_config: object
-    ) -> Sequence[str]:
-        # Admission and ORM reload finish before the legacy runtime's model I/O.
-        assert len(self.read_sessions) == 2
-        assert all(not session.in_transaction() and not session.identity_map for session in self.read_sessions)
-        assert self.tenant_ids == [tenant_id]
-        assert instruction_prompt == "Suggest concise follow-ups"
-        assert model_config is None
-        assert get_credit_usage_metadata() == {"app_type": self.app_type}
-        self.histories.append(histories)
+    def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int:
+        self.assert_queries_closed()
+        self.token_calls += 1
         if self.failure is not None:
             raise self.failure
-        return self.questions
+        return len(prompt_messages)
+
+    def get_model_schema(self) -> AIModelEntity:
+        self.assert_queries_closed()
+        return AIModelEntity.model_construct(parameter_rules=[])
+
+    def invoke_llm(
+        self,
+        *,
+        prompt_messages: list[PromptMessage],
+        model_parameters: Mapping[str, object],
+        stop: list[str],
+        stream: bool,
+    ) -> LLMResult:
+        self.assert_queries_closed()
+        assert self.tenant_ids == [self.tenant_ids[0]] * 2
+        assert get_credit_usage_metadata() == {"app_type": self.app_type, "created_by": "suggested_questions"}
+        assert model_parameters == {"max_tokens": 256, "temperature": 0.0}
+        assert stop == []
+        assert stream is False
+        prompt = prompt_messages[0].get_text_content()
+        assert "Suggest concise follow-ups" in prompt
+        self.prompts.append(prompt)
+        return LLMResult(
+            model="question-model",
+            message=AssistantPromptMessage(content=json.dumps(self.questions)),
+            usage=LLMUsage.empty_usage(),
+        )
 
     def add_trace_task(self, task: TraceTask) -> None:
         self.traces.append(task)
@@ -196,7 +226,7 @@ def harness(
     def track_read(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None:
         read_sessions.append(session)
 
-    provider = _Provider(read_sessions)
+    provider = _Provider(read_sessions, target.id)
     services = _ApplicationServices(
         trial_apps=_TrialAppServices(
             access=TrialAppAccessService(apps=TrialAppRepository(session_factory=read_factory))
@@ -228,7 +258,6 @@ def harness(
         return provider
 
     monkeypatch.setattr(message_module.ModelManager, "for_tenant", model_manager)
-    monkeypatch.setattr(message_module.LLMGenerator, "generate_suggested_questions_after_answer", provider.generate)
     monkeypatch.setattr(message_module, "TraceQueueManager", trace_manager)
 
     app = Flask(__name__)
@@ -302,13 +331,14 @@ def test_questions_keep_response_history_owner_and_usage(
     assert response.get_json() == {"data": questions}
     assert response.headers["Content-Type"] == "application/json"
     assert int(response.headers["Content-Length"]) == len(response.data)
-    assert harness.provider.histories == ["Human: What is a trial?\nAssistant: A way to try an app."]
-    assert harness.provider.tenant_ids == [harness.target.tenant_id]
+    assert len(harness.provider.prompts) == 1
+    assert "Human: What is a trial?\nAssistant: A way to try an app." in harness.provider.prompts[0]
+    assert harness.provider.tenant_ids == [harness.target.tenant_id] * 2
     assert harness.target.tenant_id not in {harness.trial.tenant_id, harness.account.current_tenant_id}
     assert len(harness.provider.traces) == 1
     assert harness.features.events == ["setup", "csrf", "feature"]
     assert harness.usage() == 1
-    assert len(harness.legacy_sessions) == 1
+    assert len(harness.provider.model_sessions) == 2
     harness.assert_closed()
 
 
@@ -341,7 +371,7 @@ def test_message_and_conversation_require_complete_owner_chain(harness: _Harness
     response = harness.get(message_id=message_id)
 
     _assert_error(response, 404, "not_found", "Message not found" if entity == "message" else "Conversation not found")
-    assert harness.provider.histories == []
+    assert harness.provider.prompts == []
     assert harness.usage() is None
     harness.assert_closed()
 
@@ -350,7 +380,7 @@ def test_deleted_conversation_is_not_visible(harness: _Harness) -> None:
     with harness.factory.begin() as session:
         session.execute(update(Conversation).where(Conversation.id == harness.conversation.id).values(is_deleted=True))
     _assert_error(harness.get(), 404, "not_found", "Conversation not found")
-    assert harness.provider.histories == []
+    assert harness.provider.prompts == []
     harness.assert_closed()
 
 
@@ -363,7 +393,7 @@ def test_disabled_or_unset_suggestions_have_specific_error(harness: _Harness, co
             .values(suggested_questions_after_answer=config)
         )
     _assert_error(harness.get(), 403, "app_suggested_questions_after_answer_disabled")
-    assert harness.provider.histories == []
+    assert harness.provider.prompts == []
     harness.assert_closed()
 
 
@@ -372,7 +402,7 @@ def test_absent_default_model_keeps_empty_success(harness: _Harness) -> None:
     response = harness.get()
     assert response.status_code == 200
     assert response.get_json() == {"data": []}
-    assert harness.provider.histories == []
+    assert harness.provider.prompts == []
     assert harness.provider.traces == []
     harness.assert_closed()
 
@@ -425,7 +455,7 @@ def test_admission_blocks_before_runtime(
     else:
         _assert_error(response, status, code)
     assert harness.legacy_sessions == []
-    assert harness.provider.histories == []
+    assert harness.provider.prompts == []
     harness.assert_closed()
 
 
@@ -442,12 +472,13 @@ def test_admission_blocks_before_runtime(
         (InvokeError("Provider unavailable"), "completion_request_error", "Provider unavailable"),
     ],
 )
-def test_provider_failures_preserve_specific_errors_and_close_session(
+def test_token_count_failures_preserve_specific_errors_and_close_session(
     harness: _Harness, failure: Exception, code: str, message: str | None
 ) -> None:
     harness.provider.failure = failure
     _assert_error(harness.get(), 400, code, message)
-    assert len(harness.provider.histories) == 1
+    assert harness.provider.token_calls == 1
+    assert harness.provider.prompts == []
     assert harness.provider.traces == []
     assert harness.usage() is None
     harness.assert_closed()
@@ -469,11 +500,11 @@ def test_runtime_preserves_outer_request_session(harness: _Harness, failure: Exc
         assert outer_session.in_transaction()
         assert app_model in outer_session.dirty
         assert app_model.name == "Pending caller change"
-        assert len(harness.legacy_sessions) == 2
-        inner_session = harness.legacy_sessions[1]
-        assert inner_session is not outer_session
-        assert not inner_session.in_transaction()
-        assert not inner_session.identity_map
+        assert harness.provider.model_sessions
+        for inner_session in harness.provider.model_sessions:
+            assert inner_session is not outer_session
+            assert not inner_session.in_transaction()
+            assert not inner_session.identity_map
     harness.assert_closed()
     with harness.factory() as session:
         saved_app = session.get(App, harness.target.id)
@@ -488,7 +519,7 @@ def test_advanced_chat_without_published_workflow_keeps_empty_success(harness: _
     assert response.status_code == 200
     assert response.get_json() == {"data": []}
     assert len(harness.provider.read_sessions) == 2
-    assert harness.provider.histories == []
+    assert harness.provider.prompts == []
     assert harness.usage() is None
     harness.assert_closed()
 
@@ -528,6 +559,6 @@ def test_reload_errors_have_explicit_http_mapping(
 
     monkeypatch.setattr(MessageSuggestedQuestionsRuntime, "get_suggested_questions", reject_reload)
     _assert_error(harness.get(), status, code, message)
-    assert harness.provider.histories == []
+    assert harness.provider.prompts == []
     assert harness.usage() is None
     harness.assert_closed()

api/tests/unit_tests/services/test_message_suggested_questions_adapters.py

@@ -7,22 +7,26 @@ from decimal import Decimal
 from typing import Literal, cast
 from uuid import uuid4
 
+import httpx
 import pytest
-from flask import Flask
+from flask import Flask, g
 from sqlalchemy import Connection, Engine, event, update
 from sqlalchemy.orm import Session, SessionTransaction, sessionmaker
 
 import services.message_service as message_module
-from core.credit_usage import CreditUsageAppType
+from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy
+from core.file import remote_fetcher
 from core.model_context import get_credit_usage_metadata, use_credit_usage_metadata
 from core.model_manager import ModelInstance
 from core.ops.ops_trace_manager import TraceTask
 from extensions.ext_database import db
-from graphon.model_runtime.entities import PromptMessage
-from graphon.model_runtime.entities.model_entities import ModelType
+from graphon.file import FileTransferMethod, FileType
+from graphon.model_runtime.entities import AssistantPromptMessage, PromptMessage
+from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
+from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from models import Account, App, AppMode, Conversation, Message
-from models.enums import ConversationFromSource, EndUserType
-from models.model import AppModelConfig, EndUser
+from models.enums import ConversationFromSource, CreatorUserRole, EndUserType
+from models.model import AppModelConfig, EndUser, MessageFile
 from models.workflow import Workflow, WorkflowType
 from services.app_definition_query_service import AppDefinitionUnavailableError
 from services.errors.conversation import ConversationNotExistsError
@@ -42,37 +46,83 @@ type _InvokeFrom = Literal["debugger", "explore", "web-app", "service-api"]
 class _Provider:
     read_sessions: list[Session]
     tenant_id: str
-    prompts: list[str | None] = field(default_factory=list)
-    histories: list[str] = field(default_factory=list)
-    model_configs: list[object] = field(default_factory=list)
+    app_id: str
+    prompts: list[str] = field(default_factory=list)
+    model_requests: list[tuple[str, str]] = field(default_factory=list)
+    invocation_parameters: list[Mapping[str, object]] = field(default_factory=list)
     metadata: list[Mapping[str, object] | None] = field(default_factory=list)
     traces: list[TraceTask] = field(default_factory=list)
+    model_sessions: list[Session] = field(default_factory=list)
+    io_sessions: list[tuple[str, bool]] = field(default_factory=list)
     failure: Exception | None = None
+    token_failure: Exception | None = None
     history_model_failure: bool = False
+    generation_model_failure: bool = False
 
-    def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance:
+    def record_io(self, stage: str) -> None:
+        self.io_sessions.append(
+            (
+                stage,
+                all(
+                    not session.in_transaction() and not session.identity_map
+                    for session in self.read_sessions + self.model_sessions
+                ),
+            )
+        )
+
+    def _resolve_model(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance:
         assert tenant_id == self.tenant_id
         assert model_type == ModelType.LLM
         assert self.read_sessions
         assert all(not session.in_transaction() and not session.identity_map for session in self.read_sessions)
-        if self.history_model_failure:
-            raise RuntimeError("No history model")
+        # Simulate ModelManager's real scoped database lookup, so both resolution
+        # stages must dispose their sessions before token/schema/provider I/O.
+        session = db.session()
+        assert session.get(App, self.app_id) is not None
+        self.model_sessions.append(session)
+        if self.history_model_failure or (self.generation_model_failure and len(self.model_sessions) > 1):
+            raise RuntimeError("No model available")
         return cast(ModelInstance, self)
 
+    def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance:
+        self.model_requests.append(("default", "default"))
+        return self._resolve_model(tenant_id=tenant_id, model_type=model_type)
+
+    def get_model_instance(self, *, tenant_id: str, model_type: ModelType, provider: str, model: str) -> ModelInstance:
+        self.model_requests.append((provider, model))
+        return self._resolve_model(tenant_id=tenant_id, model_type=model_type)
+
     def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int:
+        self.record_io("tokens")
+        if self.token_failure is not None:
+            raise self.token_failure
         return len(prompt_messages)
 
-    def generate(
-        self, *, tenant_id: str, histories: str, instruction_prompt: str | None, model_config: object
-    ) -> list[str]:
-        assert tenant_id == self.tenant_id
-        self.prompts.append(instruction_prompt)
-        self.histories.append(histories)
-        self.model_configs.append(model_config)
+    def get_model_schema(self) -> AIModelEntity:
+        self.record_io("schema")
+        return AIModelEntity.model_construct(parameter_rules=[])
+
+    def invoke_llm(
+        self,
+        *,
+        prompt_messages: list[PromptMessage],
+        model_parameters: Mapping[str, object],
+        stop: list[str],
+        stream: bool,
+    ) -> LLMResult:
+        self.record_io("invoke")
+        assert stream is False
+        assert stop == []
+        self.invocation_parameters.append(model_parameters)
+        self.prompts.append(prompt_messages[0].get_text_content())
         self.metadata.append(get_credit_usage_metadata())
         if self.failure is not None:
             raise self.failure
-        return ["What next?"]
+        return LLMResult(
+            model="question-model",
+            message=AssistantPromptMessage(content='["What next?"]'),
+            usage=LLMUsage.empty_usage(),
+        )
 
     def add_trace_task(self, task: TraceTask) -> None:
         self.traces.append(task)
@@ -121,6 +171,7 @@ class _Harness:
             not session.in_transaction() and not session.identity_map
             for session in self.provider.read_sessions + self.legacy_sessions
         )
+        assert all(closed for _stage, closed in self.provider.io_sessions), self.provider.io_sessions
 
 
 @pytest.fixture
@@ -176,7 +227,7 @@ def harness(
     def track_read(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None:
         read_sessions.append(session)
 
-    provider = _Provider(read_sessions, target.tenant_id)
+    provider = _Provider(read_sessions, target.tenant_id, target.id)
 
     def model_manager(*, tenant_id: str) -> _Provider:
         assert tenant_id == target.tenant_id
@@ -187,7 +238,6 @@ def harness(
         return provider
 
     monkeypatch.setattr(message_module.ModelManager, "for_tenant", model_manager)
-    monkeypatch.setattr(message_module.LLMGenerator, "generate_suggested_questions_after_answer", provider.generate)
     monkeypatch.setattr(message_module, "TraceQueueManager", trace_manager)
     flask_app = Flask(__name__)
     flask_app.config.update(TESTING=True, SQLALCHEMY_DATABASE_URI=str(sqlite_engine.url))
@@ -279,25 +329,32 @@ def test_runtime_preserves_history_configuration_credit_context_and_caller_sessi
         original_metadata = get_credit_usage_metadata()
         with use_credit_usage_metadata({"request_id": "suggestions-request"}):
             inherited_metadata = get_credit_usage_metadata()
-            if provider_failure:
-                with pytest.raises(RuntimeError, match="Provider unavailable"):
-                    harness.get(actor)
-            else:
-                assert harness.get(actor) == ["What next?"]
+            assert harness.get(actor) == ([] if provider_failure else ["What next?"])
             assert get_credit_usage_metadata() == inherited_metadata
-            assert harness.provider.metadata == [{**(inherited_metadata or {}), "app_type": CreditUsageAppType.CHATBOT}]
+            assert harness.provider.metadata == [
+                {
+                    **(inherited_metadata or {}),
+                    "app_type": CreditUsageAppType.CHATBOT,
+                    "created_by": CreditUsageCreatedBy.SUGGESTED_QUESTIONS,
+                }
+            ]
         assert get_credit_usage_metadata() == original_metadata
         assert db.session() is caller_session
         assert caller_session.in_transaction()
         assert caller_app in caller_session.dirty
         assert caller_app.name == "Pending caller edit"
-        assert len(harness.legacy_sessions) == 1
-        assert harness.legacy_sessions[0] is not caller_session
+        assert len(harness.provider.model_sessions) == 2
+        assert harness.provider.model_sessions[0] is not harness.provider.model_sessions[1]
+        assert all(session is not caller_session for session in harness.legacy_sessions)
         harness.assert_closed()
-    assert harness.provider.prompts == ["Published questions"]
-    assert harness.provider.histories == ["Human: How does this work?\nAssistant: Like this."]
-    assert harness.provider.model_configs == [model_config]
-    assert len(harness.provider.traces) == (0 if provider_failure else 1)
+    assert len(harness.provider.prompts) == 1
+    assert harness.provider.prompts[0].startswith(
+        "Human: How does this work?\nAssistant: Like this.\nPublished questions\n"
+    )
+    assert harness.provider.model_requests == [("default", "default"), ("vendor", "question-model")]
+    assert [stage for stage, _closed in harness.provider.io_sessions] == ["tokens", "invoke"]
+    assert harness.provider.invocation_parameters == [{}]
+    assert len(harness.provider.traces) == 1
 
 
 @pytest.mark.parametrize("invoke_from", ["explore", "service-api"])
@@ -321,7 +378,7 @@ def test_message_and_conversation_ownership_are_enforced(
     error_type = MessageNotExistsError if entity == "message" else ConversationNotExistsError
     with harness.flask_app.app_context(), pytest.raises(error_type):
         harness.get(actor)
-    assert not harness.provider.histories
+    assert not harness.provider.prompts
     harness.assert_closed()
 
 
@@ -334,7 +391,7 @@ def test_actor_cannot_read_a_deleted_conversation(
         session.execute(update(Conversation).where(Conversation.id == harness.conversation.id).values(is_deleted=True))
     with harness.flask_app.app_context(), pytest.raises(ConversationNotExistsError):
         harness.get(actor)
-    assert not harness.provider.histories
+    assert not harness.provider.prompts
     harness.assert_closed()
 
 
@@ -362,8 +419,8 @@ def test_account_preserves_disabled_feature_and_missing_workflow_or_model_behavi
                 harness.get(actor)
         else:
             assert harness.get(actor) == []
-    assert len(harness.legacy_sessions) == 1
-    assert not harness.provider.histories
+    assert len(harness.legacy_sessions) == (1 if case == "history_provider_failure" else 0)
+    assert not harness.provider.prompts
     assert not harness.provider.traces
     harness.assert_closed()
 
@@ -400,6 +457,181 @@ def test_advanced_chat_debugger_uses_draft_while_explore_uses_published_workflow
             actor=actor,
             message_id=harness.message.id,
         ) == ["What next?"]
-    assert harness.provider.prompts == ["Draft questions" if invoke_from == "debugger" else "Published questions"]
-    assert harness.provider.model_configs == [None]
+    expected_prompt = "Draft questions" if invoke_from == "debugger" else "Published questions"
+    assert len(harness.provider.prompts) == 1
+    assert expected_prompt in harness.provider.prompts[0]
+    assert harness.provider.model_requests == [("default", "default"), ("default", "default")]
+    assert [stage for stage, _closed in harness.provider.io_sessions] == ["tokens", "schema", "invoke"]
     harness.assert_closed()
+
+
+@pytest.mark.parametrize("failure_stage", ["tokens", "generation_model"])
+def test_runtime_releases_sessions_when_history_or_second_model_stage_fails(
+    harness: _Harness, failure_stage: Literal["tokens", "generation_model"]
+) -> None:
+    actor = harness.actor("explore")
+    if failure_stage == "tokens":
+        harness.provider.token_failure = RuntimeError("Token counter unavailable")
+    else:
+        harness.provider.generation_model_failure = True
+    with harness.flask_app.app_context():
+        if failure_stage == "tokens":
+            with pytest.raises(RuntimeError, match="Token counter unavailable"):
+                harness.get(actor)
+        else:
+            assert harness.get(actor) == []
+    assert len(harness.provider.model_sessions) == (1 if failure_stage == "tokens" else 2)
+    assert not harness.provider.prompts
+    assert len(harness.provider.traces) == (0 if failure_stage == "tokens" else 1)
+    harness.assert_closed()
+
+
+@pytest.mark.parametrize("remote_failure", [False, True])
+def test_remote_attachment_is_loaded_after_history_and_model_sessions_close(
+    harness: _Harness, monkeypatch: pytest.MonkeyPatch, remote_failure: bool
+) -> None:
+    actor = harness.actor("explore")
+    with harness.factory.begin() as session:
+        session.execute(
+            update(AppModelConfig)
+            .where(AppModelConfig.id == harness.target.app_model_config_id)
+            .values(
+                file_upload=json.dumps(
+                    {"enabled": True, "allowed_file_types": ["image"], "allowed_file_upload_methods": ["remote_url"]}
+                )
+            )
+        )
+        session.add(
+            MessageFile(
+                message_id=harness.message.id,
+                type=FileType.IMAGE,
+                transfer_method=FileTransferMethod.REMOTE_URL,
+                url="https://example.com/history.png",
+                created_by_role=CreatorUserRole.ACCOUNT,
+                created_by=harness.account.id,
+            )
+        )
+    requests: list[str] = []
+
+    def fetch(method: str, url: str, **_kwargs: object) -> httpx.Response:
+        harness.provider.record_io("attachment")
+        requests.append(method)
+        assert url == "https://example.com/history.png"
+        if remote_failure:
+            raise RuntimeError("Attachment unavailable")
+        return httpx.Response(
+            200,
+            headers={"content-type": "image/png", "content-length": "3"},
+            content=b"png",
+            request=httpx.Request(method, url),
+        )
+
+    monkeypatch.setattr(remote_fetcher, "make_request", fetch)
+    with harness.flask_app.app_context():
+        if remote_failure:
+            with pytest.raises(RuntimeError, match="Attachment unavailable"):
+                harness.get(actor)
+        else:
+            assert harness.get(actor) == ["What next?"]
+    assert requests
+    assert requests[0] == "HEAD"
+    if remote_failure:
+        assert not harness.provider.prompts
+        assert not harness.provider.traces
+    else:
+        assert harness.provider.prompts[0].startswith("Human: [image]\nHow does this work?\nAssistant: Like this.")
+    harness.assert_closed()
+
+
+@pytest.mark.parametrize("failure_stage", [None, "provider_sql", "trace"])
+def test_runtime_keeps_its_context_state_while_releasing_sessions_and_preserving_the_caller(
+    harness: _Harness, monkeypatch: pytest.MonkeyPatch, failure_stage: Literal["provider_sql", "trace"] | None
+) -> None:
+    actor = harness.actor("explore")
+    runtime_steps: list[str] = []
+    invocation_sessions: list[Session] = []
+    trace_sessions: list[Session] = []
+    resolve_model = harness.provider._resolve_model
+    count_tokens = harness.provider.get_llm_num_tokens
+    invoke_llm = harness.provider.invoke_llm
+
+    def resolve(*, tenant_id: str, model_type: ModelType) -> ModelInstance:
+        assert g.get("caller_marker") is None
+        if not runtime_steps:
+            g.suggested_question_steps = runtime_steps
+        assert g.get("suggested_question_steps") is runtime_steps
+        runtime_steps.append("model")
+        return resolve_model(tenant_id=tenant_id, model_type=model_type)
+
+    def tokens(prompt_messages: Sequence[PromptMessage]) -> int:
+        assert g.get("suggested_question_steps") is runtime_steps
+        runtime_steps.append("tokens")
+        return count_tokens(prompt_messages)
+
+    def invoke(
+        *, prompt_messages: list[PromptMessage], model_parameters: Mapping[str, object], stop: list[str], stream: bool
+    ) -> LLMResult:
+        assert g.get("suggested_question_steps") is runtime_steps
+        runtime_steps.append("invoke")
+        if failure_stage == "provider_sql":
+            # A real failed flush poisons this scoped session. LLMGenerator
+            # catches the error, so tracing must receive a fresh session.
+            session = db.session()
+            invocation_sessions.append(session)
+            session.add(
+                App(id=harness.target.id, tenant_id=harness.target.tenant_id, name="Duplicate", mode=AppMode.CHAT)
+            )
+            try:
+                session.flush()
+            finally:
+                assert not session.is_active
+        return invoke_llm(prompt_messages=prompt_messages, model_parameters=model_parameters, stop=stop, stream=stream)
+
+    def trace(task: TraceTask) -> None:
+        assert g.get("suggested_question_steps") is runtime_steps
+        runtime_steps.append("trace")
+        session = db.session()
+        trace_sessions.append(session)
+        app = session.get(App, harness.target.id)
+        assert app is not None
+        assert app.name == "Shared"
+        if failure_stage == "trace":
+            raise RuntimeError("Trace backend unavailable")
+        harness.provider.traces.append(task)
+
+    monkeypatch.setattr(harness.provider, "_resolve_model", resolve)
+    monkeypatch.setattr(harness.provider, "get_llm_num_tokens", tokens)
+    monkeypatch.setattr(harness.provider, "invoke_llm", invoke)
+    monkeypatch.setattr(harness.provider, "add_trace_task", trace)
+    with harness.flask_app.app_context():
+        caller_marker = object()
+        g.caller_marker = caller_marker
+        caller_session = db.session()
+        caller_app = caller_session.get(App, harness.target.id)
+        assert caller_app is not None
+        caller_app.name = "Pending caller edit"
+        harness.legacy_sessions.clear()
+        if failure_stage == "trace":
+            with pytest.raises(RuntimeError, match="Trace backend unavailable"):
+                harness.get(actor)
+        else:
+            assert harness.get(actor) == ([] if failure_stage == "provider_sql" else ["What next?"])
+        assert g.caller_marker is caller_marker
+        assert g.get("suggested_question_steps") is None
+        assert db.session() is caller_session
+        assert caller_session.in_transaction()
+        assert caller_app in caller_session.dirty
+        assert caller_app.name == "Pending caller edit"
+        assert all(session is not caller_session for session in harness.legacy_sessions)
+        harness.assert_closed()
+
+    assert runtime_steps == ["model", "tokens", "model", "invoke", "trace"]
+    assert len(trace_sessions) == 1
+    assert not trace_sessions[0].in_transaction()
+    if failure_stage == "provider_sql":
+        assert len(invocation_sessions) == 1
+        assert trace_sessions[0] is not invocation_sessions[0]
+    with harness.factory() as session:
+        persisted_app = session.get(App, harness.target.id)
+        assert persisted_app is not None
+        assert persisted_app.name == "Shared"