feat(tools): support tool-generated A2UI parts

Stream, validate, persist, and render constrained UI parts from tools, including the current-time component. Preserve markdown following think blocks as normal answer content.

api/configs/feature/__init__.py

@@ -1497,11 +1497,6 @@ class LoginConfig(BaseSettings):
 
 
 class AccountConfig(BaseSettings):
-    ENABLE_CHANGE_EMAIL: bool = Field(
-        description="whether users can change their email address",
-        default=True,
-    )
-
     ACCOUNT_DELETION_TOKEN_EXPIRY_MINUTES: PositiveInt = Field(
         description="Duration in minutes for which a account deletion token remains valid",
         default=5,

api/controllers/console/agent/roster.py

@@ -539,9 +539,23 @@ def _parse_observability_time_range(start: str | None, end: str | None, account:
 
 
 def _query_values(name: str, alias_name: str | None = None) -> list[str]:
-    values = request.args.getlist(name)
+    def _get_values(field_name: str) -> list[str]:
+        values = request.args.getlist(field_name)
+        indexed_values: list[tuple[int, list[str]]] = []
+        prefix = f"{field_name}["
+        for key in request.args:
+            if not key.startswith(prefix) or not key.endswith("]"):
+                continue
+            index = key[len(prefix) : -1]
+            if index.isdigit():
+                indexed_values.append((int(index), request.args.getlist(key)))
+        for _, items in sorted(indexed_values):
+            values.extend(items)
+        return values
+
+    values = _get_values(name)
     if alias_name:
-        values.extend(request.args.getlist(alias_name))
+        values.extend(_get_values(alias_name))
     return [value.strip() for value in values if value.strip()]
 
 

api/controllers/console/app/completion.py

@@ -351,24 +351,31 @@ def _resolve_current_user_agent_debug_conversation_id(
     app_model: App,
     agent_id: str | None,
     draft_type: AgentConfigDraftType,
+    start_new: bool = False,
 ) -> str:
-    """Resolve the current editor's conversation without crossing draft surfaces."""
+    """Resolve or rotate the current editor's conversation within one draft surface.
+
+    ``start_new`` rotates the scoped mapping through ``AgentRosterService`` so
+    the old runtime session is retired before the new conversation is used.
+    Continuations and Build chat keep resolving the existing mapping.
+    """
 
     roster_service = AgentRosterService(session)
-    if agent_id:
-        return roster_service.get_or_create_agent_app_debug_conversation_id(
-            tenant_id=current_tenant_id,
-            agent_id=agent_id,
-            account_id=current_user.id,
-            draft_type=draft_type,
-        )
+    resolved_agent_id = agent_id
+    if not resolved_agent_id:
+        agent = roster_service.get_app_backing_agent(tenant_id=current_tenant_id, app_id=str(app_model.id))
+        if agent is None:
+            raise AgentNotFoundError()
+        resolved_agent_id = agent.id
 
-    agent = roster_service.get_app_backing_agent(tenant_id=current_tenant_id, app_id=str(app_model.id))
-    if agent is None:
-        raise AgentNotFoundError()
-    return roster_service.get_or_create_agent_app_debug_conversation_id(
+    resolve_conversation = (
+        roster_service.refresh_agent_app_debug_conversation_id
+        if start_new
+        else roster_service.get_or_create_agent_app_debug_conversation_id
+    )
+    return resolve_conversation(
         tenant_id=current_tenant_id,
-        agent_id=agent.id,
+        agent_id=resolved_agent_id,
         account_id=current_user.id,
         draft_type=draft_type,
     )
@@ -387,13 +394,17 @@ def _create_chat_message(
     args = args_model.model_dump(exclude_none=True, by_alias=True)
 
     if AppMode.value_of(app_model.mode) == AppMode.AGENT:
+        draft_type = AgentConfigDraftType(args_model.draft_type)
+        # Preview follows the normal chat contract: an omitted/empty conversation ID starts a new
+        # conversation. Build chat keeps its stable mapping so build drafts and finalization stay continuous.
         debug_conversation_id = _resolve_current_user_agent_debug_conversation_id(
             session=session,
             current_tenant_id=current_tenant_id or app_model.tenant_id,
             current_user=current_user,
             app_model=app_model,
             agent_id=agent_id,
-            draft_type=AgentConfigDraftType(args_model.draft_type),
+            draft_type=draft_type,
+            start_new=draft_type == AgentConfigDraftType.DRAFT and not args_model.conversation_id,
         )
         if args_model.conversation_id and args_model.conversation_id != debug_conversation_id:
             raise NotFound("Conversation Not Exists.")

api/controllers/console/auth/oauth_server.py

@@ -56,6 +56,7 @@ class OAuthProviderTokenResponse(BaseModel):
 
 
 class OAuthProviderAccountResponse(BaseModel):
+    id: str
     name: str
     email: str
     avatar: str | None = None
@@ -251,6 +252,7 @@ class OAuthServerUserAccountApi(Resource):
     def post(self, oauth_provider_app: OAuthProviderApp, account: Account):
         return jsonable_encoder(
             {
+                "id": account.id,
                 "name": account.name,
                 "email": account.email,
                 "avatar": account.avatar,

api/controllers/web/wraps.py

@@ -11,6 +11,7 @@ from werkzeug.exceptions import BadRequest, NotFound, Unauthorized
 
 from constants import HEADER_NAME_APP_CODE
 from controllers.web.error import WebAppAuthAccessDeniedError, WebAppAuthRequiredError
+from core.logging.context import set_identity_context
 from extensions.ext_database import db
 from libs.passport import PassportService
 from libs.token import extract_webapp_passport
@@ -28,6 +29,11 @@ def validate_jwt_token[**P, R](
         @wraps(view)
         def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
             app_model, end_user = decode_jwt_token()
+            set_identity_context(
+                tenant_id=end_user.tenant_id,
+                user_id=end_user.id,
+                user_type=end_user.type or "end_user",
+            )
             return view(app_model, end_user, *args, **kwargs)
 
         return decorated

api/core/agent/base_agent_runner.py

@@ -10,12 +10,13 @@ from sqlalchemy.orm import Session
 from core.agent.entities import AgentEntity, AgentToolEntity
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
 from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig
-from core.app.apps.base_app_queue_manager import AppQueueManager
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.apps.base_app_runner import AppRunner
 from core.app.entities.app_invoke_entities import (
     AgentChatAppGenerateEntity,
     ModelConfigWithCredentialsEntity,
 )
+from core.app.entities.queue_entities import QueueUIPartEvent
 from core.app.file_access import DatabaseFileAccessController
 from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
@@ -23,6 +24,12 @@ from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance
 from core.prompt.utils.extract_thread_messages import extract_thread_messages
 from core.tools.__base.tool import Tool
+from core.tools.entities.ui_entities import (
+    MessageUIPart,
+    ToolUIMessage,
+    build_ui_part_id,
+    validate_tool_ui_message_batch,
+)
 from core.tools.tool_manager import ToolManager
 from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool
 from extensions.ext_database import db
@@ -352,6 +359,34 @@ class BaseAgentRunner(AppRunner):
         db.session.commit()
         db.session.close()
 
+    def publish_ui_messages(self, *, namespace: str, ui_messages: list[ToolUIMessage]) -> None:
+        """Publish validated tool UI messages without adding them to agent observations."""
+        try:
+            validate_tool_ui_message_batch(ui_messages)
+        except ValueError:
+            logger.warning(
+                "Ignored tool UI batch that exceeds queue publication limits",
+                extra={"message_id": self.message.id, "namespace": namespace},
+                exc_info=True,
+            )
+            return
+
+        sequences: dict[str, int] = {}
+        for ui_message in ui_messages:
+            part_id = build_ui_part_id(namespace, ui_message.surface_id)
+            sequence = sequences.get(part_id, 0) + 1
+            sequences[part_id] = sequence
+            self.queue_manager.publish(
+                QueueUIPartEvent(
+                    part=MessageUIPart.from_tool_ui_message(
+                        part_id=part_id,
+                        sequence=sequence,
+                        ui_message=ui_message,
+                    )
+                ),
+                PublishFrom.APPLICATION_MANAGER,
+            )
+
     def organize_agent_history(self, prompt_messages: list[PromptMessage], *, session: Session) -> list[PromptMessage]:
         """
         Organize agent history

api/core/agent/cot_agent_runner.py

@@ -235,6 +235,7 @@ class CotAgentRunner(BaseAgentRunner, ABC):
                     # action is tool call, invoke tool
                     tool_invoke_response, tool_invoke_meta = self._handle_invoke_action(
                         session=session,
+                        agent_thought_id=agent_thought_id,
                         action=scratchpad.action,
                         tool_instances=tool_instances,
                         message_file_ids=message_file_ids,
@@ -306,9 +307,11 @@ class CotAgentRunner(BaseAgentRunner, ABC):
         tool_instances: Mapping[str, Tool],
         message_file_ids: list[str],
         trace_manager: TraceQueueManager | None = None,
+        agent_thought_id: str | None = None,
     ) -> tuple[str, ToolInvokeMeta]:
         """
         handle invoke action
+        :param agent_thought_id: namespace for UI parts emitted by this action
         :param action: action
         :param tool_instances: tool instances
         :param message_file_ids: message file ids
@@ -331,7 +334,7 @@ class CotAgentRunner(BaseAgentRunner, ABC):
                 pass
 
         # invoke tool
-        tool_invoke_response, message_files, tool_invoke_meta = ToolEngine.agent_invoke(
+        invoke_result = ToolEngine.agent_invoke(
             session=session,
             tool=tool_instance,
             tool_parameters=tool_call_args,
@@ -342,6 +345,13 @@ class CotAgentRunner(BaseAgentRunner, ABC):
             agent_tool_callback=self.agent_callback,
             trace_manager=trace_manager,
         )
+        tool_invoke_response = invoke_result.observation
+        message_files = invoke_result.message_files
+        tool_invoke_meta = invoke_result.meta
+        self.publish_ui_messages(
+            namespace=agent_thought_id or self.message.id,
+            ui_messages=invoke_result.ui_messages,
+        )
         session.commit()
         session.close()
 

api/core/agent/fc_agent_runner.py

@@ -251,7 +251,7 @@ class FunctionCallAgentRunner(BaseAgentRunner):
                     }
                 else:
                     # invoke tool
-                    tool_invoke_response, message_files, tool_invoke_meta = ToolEngine.agent_invoke(
+                    invoke_result = ToolEngine.agent_invoke(
                         session=session,
                         tool=tool_instance,
                         tool_parameters=tool_call_args,
@@ -265,6 +265,13 @@ class FunctionCallAgentRunner(BaseAgentRunner):
                         message_id=self.message.id,
                         conversation_id=self.conversation.id,
                     )
+                    tool_invoke_response = invoke_result.observation
+                    message_files = invoke_result.message_files
+                    tool_invoke_meta = invoke_result.meta
+                    self.publish_ui_messages(
+                        namespace=tool_call_id,
+                        ui_messages=invoke_result.ui_messages,
+                    )
                     session.commit()
                     session.close()
                     # publish files

api/core/app/apps/agent_app/app_runner.py

@@ -52,8 +52,16 @@ from core.app.entities.queue_entities import (
     QueueAgentThoughtEvent,
     QueueLLMChunkEvent,
     QueueMessageEndEvent,
+    QueueUIPartEvent,
 )
 from core.repositories.human_input_repository import HumanInputFormRepository, HumanInputFormRepositoryImpl
+from core.tools.entities.ui_entities import (
+    MessageUIPart,
+    ToolUIMessage,
+    build_ui_part_id,
+    parse_tool_ui_messages,
+    validate_tool_ui_message_batch,
+)
 from core.workflow.nodes.agent_v2.ask_human_hitl import AskHumanFormBuildError, create_ask_human_form
 from core.workflow.nodes.agent_v2.ask_human_resume import build_deferred_tool_results, resolve_ask_human_form
 from extensions.ext_database import db
@@ -435,9 +443,29 @@ class _AgentProcessRecorder:
         content = part.get("content")
         if content is None:
             content = part
-        self._record_tool_observation(tool_call_id=tool_call_id, tool_name=tool_name, observation=content)
+        thought_id = self._record_tool_observation(
+            tool_call_id=tool_call_id,
+            tool_name=tool_name,
+            observation=content,
+        )
 
-    def _record_tool_observation(self, *, tool_call_id: str | None, tool_name: str | None, observation: Any) -> None:
+        metadata = part.get("metadata")
+        if not isinstance(metadata, Mapping) or "dify_ui_messages" not in metadata:
+            return
+        try:
+            ui_messages = parse_tool_ui_messages(metadata["dify_ui_messages"])
+        except (TypeError, ValueError):
+            logger.warning("Ignored invalid dify_ui_messages metadata from Agent backend", exc_info=True)
+            return
+        self._publish_ui_messages(namespace=tool_call_id or thought_id, ui_messages=ui_messages)
+
+    def _record_tool_observation(
+        self,
+        *,
+        tool_call_id: str | None,
+        tool_name: str | None,
+        observation: Any,
+    ) -> str:
         self._close_thinking_segments()
         thought_id = self._lookup_observation_thought(tool_call_id=tool_call_id, tool_name=tool_name)
         if thought_id is None:
@@ -445,6 +473,34 @@ class _AgentProcessRecorder:
         else:
             self._mark_tool_observed(thought_id)
         self._update_thought(thought_id, observation=_json_or_text(observation))
+        return thought_id
+
+    def _publish_ui_messages(self, *, namespace: str, ui_messages: list[ToolUIMessage]) -> None:
+        try:
+            validate_tool_ui_message_batch(ui_messages)
+        except ValueError:
+            logger.warning(
+                "Ignored Agent backend UI batch that exceeds queue publication limits",
+                extra={"message_id": self._message_id, "namespace": namespace},
+                exc_info=True,
+            )
+            return
+
+        sequences: dict[str, int] = {}
+        for ui_message in ui_messages:
+            part_id = build_ui_part_id(namespace, ui_message.surface_id)
+            sequence = sequences.get(part_id, 0) + 1
+            sequences[part_id] = sequence
+            self._queue_manager.publish(
+                QueueUIPartEvent(
+                    part=MessageUIPart.from_tool_ui_message(
+                        part_id=part_id,
+                        sequence=sequence,
+                        ui_message=ui_message,
+                    )
+                ),
+                PublishFrom.APPLICATION_MANAGER,
+            )
 
     def _lookup_tool_thought(self, *, index: int, tool_call_id: str | None) -> str | None:
         if tool_call_id and tool_call_id in self._tool_by_call_id:

api/core/app/entities/queue_entities.py

@@ -7,6 +7,7 @@ from pydantic import BaseModel, ConfigDict, Field
 
 from core.app.entities.agent_strategy import AgentStrategyInfo
 from core.rag.entities import RetrievalSourceMetadata
+from core.tools.entities.ui_entities import MessageUIPart
 from core.workflow.nodes.human_input.pause_reason import PauseReason
 from graphon.entities import WorkflowStartReason
 from graphon.enums import NodeType, WorkflowNodeExecutionMetadataKey
@@ -43,6 +44,7 @@ class QueueEvent(StrEnum):
     REASONING_CHUNK = "reasoning_chunk"
     ANNOTATION_REPLY = "annotation_reply"
     AGENT_THOUGHT = "agent_thought"
+    UI_PART = "ui_part"
     MESSAGE_FILE = "message_file"
     AGENT_LOG = "agent_log"
     ERROR = "error"
@@ -457,6 +459,13 @@ class QueueAgentThoughtEvent(AppQueueEvent):
     agent_thought_id: str
 
 
+class QueueUIPartEvent(AppQueueEvent):
+    """A validated tool UI surface revision."""
+
+    event: QueueEvent = QueueEvent.UI_PART
+    part: MessageUIPart
+
+
 class QueueMessageFileEvent(AppQueueEvent):
     """
     QueueAgentThoughtEvent entity

api/core/app/entities/task_entities.py

@@ -2,10 +2,11 @@ from collections.abc import Mapping, Sequence
 from enum import StrEnum
 from typing import Any, Literal
 
-from pydantic import BaseModel, ConfigDict, Field, JsonValue
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator
 
 from core.app.entities.agent_strategy import AgentStrategyInfo
 from core.rag.entities import RetrievalSourceMetadata
+from core.tools.entities.ui_entities import MessageUIPart, validate_ui_part_batch
 from core.workflow.nodes.human_input.entities import FormInputConfig, UserActionConfig
 from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType
 from graphon.entities import WorkflowStartReason
@@ -30,6 +31,14 @@ class TaskStateMetadata(BaseModel):
     reasoning: dict[str, str] = Field(default_factory=dict)
     """reasoning_content per LLM node id (separated mode), accumulated across iteration/loop
     passes for that node; persisted to message_metadata"""
+    ui_parts: list[MessageUIPart] = Field(default_factory=list)
+    """Latest validated revision of each tool UI part; persisted to message metadata."""
+
+    @field_validator("ui_parts")
+    @classmethod
+    def _validate_ui_parts(cls, value: list[MessageUIPart]) -> list[MessageUIPart]:
+        validate_ui_part_batch(value)
+        return value
 
 
 class TaskState(BaseModel):
@@ -73,6 +82,7 @@ class StreamEvent(StrEnum):
     MESSAGE_FILE = "message_file"
     MESSAGE_REPLACE = "message_replace"
     AGENT_THOUGHT = "agent_thought"
+    UI_PART = "ui_part"
     AGENT_MESSAGE = "agent_message"
     WORKFLOW_STARTED = "workflow_started"
     WORKFLOW_PAUSED = "workflow_paused"
@@ -192,6 +202,14 @@ class AgentThoughtStreamResponse(StreamResponse):
     message_files: list[str] | None = None
 
 
+class UIPartStreamResponse(StreamResponse):
+    """A tool-owned UI surface revision."""
+
+    event: StreamEvent = StreamEvent.UI_PART
+    id: str
+    part: MessageUIPart
+
+
 class AgentMessageStreamResponse(StreamResponse):
     """
     AgentMessageStreamResponse entity

api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py

@@ -26,6 +26,7 @@ from core.app.entities.queue_entities import (
     QueuePingEvent,
     QueueRetrieverResourcesEvent,
     QueueStopEvent,
+    QueueUIPartEvent,
 )
 from core.app.entities.task_entities import (
     AgentMessageStreamResponse,
@@ -41,6 +42,7 @@ from core.app.entities.task_entities import (
     MessageEndStreamResponse,
     StreamEvent,
     StreamResponse,
+    UIPartStreamResponse,
 )
 from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
 from core.app.task_pipeline.message_cycle_manager import MessageCycleManager
@@ -52,6 +54,7 @@ from core.ops.entities.trace_entity import TraceTaskName
 from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
 from core.prompt.utils.prompt_message_util import PromptMessageUtil
 from core.prompt.utils.prompt_template_parser import PromptTemplateParser
+from core.tools.entities.ui_entities import upsert_ui_part
 from events.message_event import message_was_created
 from graphon.file import FileTransferMethod
 from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
@@ -311,6 +314,10 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat
                     agent_thought_response = self._agent_thought_to_stream_response(event)
                     if agent_thought_response is not None:
                         yield agent_thought_response
+                case QueueUIPartEvent():
+                    ui_part_response = self._ui_part_to_stream_response(event)
+                    if ui_part_response is not None:
+                        yield ui_part_response
                 case QueueMessageFileEvent():
                     response = self._message_cycle_manager.message_file_to_stream_response(event)
                     if response:
@@ -526,6 +533,27 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat
             task_id=self._application_generate_entity.task_id, id=message_id, answer=answer
         )
 
+    def _ui_part_to_stream_response(self, event: QueueUIPartEvent) -> UIPartStreamResponse | None:
+        """Upsert a UI part revision into persistent metadata and stream it."""
+        try:
+            updated_parts = upsert_ui_part(self._task_state.metadata.ui_parts, event.part)
+        except ValueError:
+            logger.warning(
+                "Ignored UI part that exceeds assistant message limits",
+                extra={"message_id": self._message_id, "part_id": event.part.part_id},
+                exc_info=True,
+            )
+            return None
+        if updated_parts is None:
+            return None
+        self._task_state.metadata.ui_parts = updated_parts
+
+        return UIPartStreamResponse(
+            task_id=self._application_generate_entity.task_id,
+            id=self._message_id,
+            part=event.part,
+        )
+
     def _agent_thought_to_stream_response(self, event: QueueAgentThoughtEvent) -> AgentThoughtStreamResponse | None:
         """
         Agent thought to stream response.

api/core/indexing_runner.py

@@ -21,6 +21,7 @@ from core.model_manager import ModelInstance, ModelManager
 from core.rag.cleaner.clean_processor import CleanProcessor
 from core.rag.datasource.keyword.keyword_factory import Keyword
 from core.rag.docstore.dataset_docstore import DatasetDocumentStore
+from core.rag.embedding.token_counter import calculate_segment_token_counts
 from core.rag.extractor.entity.datasource_type import DatasourceType
 from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
@@ -113,17 +114,25 @@ class IndexingRunner:
                     current_user=current_user,
                     session=session,
                 )
+                token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
+                total_tokens = sum(token_counts)
                 # save segment
-                self._load_segments(dataset, requeried_document, documents, session)
+                self._load_segments(
+                    session=session,
+                    dataset=dataset,
+                    dataset_document=requeried_document,
+                    documents=documents,
+                    token_counts=token_counts,
+                )
                 session.commit()
 
                 # load
                 self._load(
-                    index_processor=index_processor,
+                    session=session,
                     dataset=dataset,
                     dataset_document=requeried_document,
                     documents=documents,
-                    session=session,
+                    total_tokens=total_tokens,
                 )
             except DocumentIsPausedError:
                 raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
@@ -190,17 +199,25 @@ class IndexingRunner:
                 current_user=current_user,
                 session=session,
             )
+            token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
+            total_tokens = sum(token_counts)
             # save segment
-            self._load_segments(dataset, requeried_document, documents, session)
+            self._load_segments(
+                session=session,
+                dataset=dataset,
+                dataset_document=requeried_document,
+                documents=documents,
+                token_counts=token_counts,
+            )
             session.commit()
 
             # load
             self._load(
-                index_processor=index_processor,
+                session=session,
                 dataset=dataset,
                 dataset_document=requeried_document,
                 documents=documents,
-                session=session,
+                total_tokens=total_tokens,
             )
         except DocumentIsPausedError:
             raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
@@ -225,7 +242,7 @@ class IndexingRunner:
             if not dataset:
                 raise ValueError("no dataset found")
 
-            # get exist document_segment list and delete
+            # get existing document segments
             document_segments = session.scalars(
                 select(DocumentSegment).where(
                     DocumentSegment.dataset_id == dataset.id,
@@ -264,15 +281,15 @@ class IndexingRunner:
                                     child_documents.append(child_document)
                                 document.children = child_documents
                         documents.append(document)
+            # Preserve the full document total even when only incomplete segments are re-indexed.
+            total_tokens = sum(document_segment.tokens for document_segment in document_segments)
             # build index
-            index_type = requeried_document.doc_form
-            index_processor = IndexProcessorFactory(index_type).init_index_processor()
             self._load(
-                index_processor=index_processor,
+                session=session,
                 dataset=dataset,
                 dataset_document=requeried_document,
                 documents=documents,
-                session=session,
+                total_tokens=total_tokens,
             )
         except DocumentIsPausedError:
             raise DocumentIsPausedError(f"Document paused, document id: {document_id}")
@@ -601,28 +618,16 @@ class IndexingRunner:
 
     def _load(
         self,
-        index_processor: BaseIndexProcessor,
+        session: Session,
         dataset: Dataset,
         dataset_document: DatasetDocument,
         documents: list[Document],
-        session: Session,
-    ):
-        """
-        insert index and update document/segment status to completed
-        """
+        total_tokens: int,
+    ) -> None:
+        """Build indexes and mark the document complete using the token total computed before hash sharding."""
 
-        embedding_model_instance = None
-        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
-            embedding_model_instance = self._get_model_manager(dataset.tenant_id).get_model_instance(
-                tenant_id=dataset.tenant_id,
-                provider=dataset.embedding_model_provider,
-                model_type=ModelType.TEXT_EMBEDDING,
-                model=dataset.embedding_model,
-            )
-
-        # chunk nodes by chunk size
+        # Build indexes using the existing hash-based worker groups.
         indexing_start_at = time.perf_counter()
-        tokens = 0
         create_keyword_thread = None
         if (
             dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
@@ -659,12 +664,11 @@ class IndexingRunner:
                             chunk_documents,
                             dataset.id,
                             dataset_document.id,
-                            embedding_model_instance,
                         )
                     )
 
                 for future in futures:
-                    tokens += future.result()
+                    future.result()
         if (
             dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
             and dataset.indexing_technique == IndexTechniqueType.ECONOMY
@@ -679,7 +683,7 @@ class IndexingRunner:
             document_id=dataset_document.id,
             after_indexing_status=IndexingStatus.COMPLETED,
             extra_update_params={
-                DatasetDocument.tokens: tokens,
+                DatasetDocument.tokens: total_tokens,
                 DatasetDocument.completed_at: naive_utc_now(),
                 DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
                 DatasetDocument.error: None,
@@ -720,8 +724,7 @@ class IndexingRunner:
         chunk_documents: list[Document],
         dataset_id: str,
         dataset_document_id: str,
-        embedding_model_instance: ModelInstance | None,
-    ):
+    ) -> None:
         with flask_app.app_context():
             with session_factory.create_session() as session:
                 dataset = session.get(Dataset, dataset_id)
@@ -735,11 +738,6 @@ class IndexingRunner:
                 # check document is paused
                 self._check_document_paused_status(dataset_document.id)
 
-                tokens = 0
-                if embedding_model_instance:
-                    page_content_list = [document.page_content for document in chunk_documents]
-                    tokens += sum(embedding_model_instance.get_text_embedding_num_tokens(page_content_list))
-
                 multimodal_documents = []
                 for document in chunk_documents:
                     if document.attachments and dataset.is_multimodal:
@@ -773,8 +771,6 @@ class IndexingRunner:
 
                 session.commit()
 
-                return tokens
-
     @staticmethod
     def _check_document_paused_status(document_id: str):
         indexing_cache_key = f"document_{document_id}_is_paused"
@@ -864,8 +860,14 @@ class IndexingRunner:
         return documents
 
     def _load_segments(
-        self, dataset: Dataset, dataset_document: DatasetDocument, documents: list[Document], session: Session
-    ):
+        self,
+        session: Session,
+        dataset: Dataset,
+        dataset_document: DatasetDocument,
+        documents: list[Document],
+        token_counts: list[int],
+    ) -> None:
+        """Persist transformed documents and their precomputed token counts before indexing starts."""
         # save node to document segment
         doc_store = DatasetDocumentStore(
             dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id
@@ -873,9 +875,10 @@ class IndexingRunner:
 
         # add document segments
         doc_store.add_documents(
+            session=session,
             docs=documents,
             save_child=dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX,
-            session=session,
+            token_counts=token_counts,
         )
 
         # update document status to indexing
@@ -900,7 +903,6 @@ class IndexingRunner:
                 DocumentSegment.indexing_at: naive_utc_now(),
             },
         )
-        pass
 
 
 class DocumentIsPausedError(Exception):

api/core/logging/context.py

@@ -6,9 +6,21 @@ using Python's contextvars for thread-safe and async-safe storage.
 
 import uuid
 from contextvars import ContextVar
+from typing import NamedTuple
+
+
+class IdentityContext(NamedTuple):
+    """Immutable identity values captured for logging."""
+
+    tenant_id: str
+    user_id: str
+    user_type: str
+
 
 _request_id: ContextVar[str] = ContextVar("log_request_id", default="")
 _trace_id: ContextVar[str] = ContextVar("log_trace_id", default="")
+_EMPTY_IDENTITY_CONTEXT = IdentityContext(tenant_id="", user_id="", user_type="")
+_identity: ContextVar[IdentityContext] = ContextVar("log_identity", default=_EMPTY_IDENTITY_CONTEXT)
 
 
 def get_request_id() -> str:
@@ -21,15 +33,35 @@ def get_trace_id() -> str:
     return _trace_id.get()
 
 
+def get_identity_context() -> IdentityContext:
+    """Get the immutable tenant, user, and user-type snapshot for logging."""
+    return _identity.get()
+
+
+def set_identity_context(
+    *, tenant_id: str | None = None, user_id: str | None = None, user_type: str | None = None
+) -> None:
+    """Set primitive identity values already resolved by an authentication boundary."""
+    _identity.set(
+        IdentityContext(
+            tenant_id=tenant_id or "",
+            user_id=user_id or "",
+            user_type=user_type or "",
+        )
+    )
+
+
 def init_request_context() -> None:
-    """Initialize request context. Call at start of each request."""
+    """Initialize request context and discard identity left by earlier work."""
     req_id = uuid.uuid4().hex[:10]
     trace_id = uuid.uuid5(uuid.NAMESPACE_DNS, req_id).hex
     _request_id.set(req_id)
     _trace_id.set(trace_id)
+    _identity.set(_EMPTY_IDENTITY_CONTEXT)
 
 
 def clear_request_context() -> None:
-    """Clear request context. Call at end of request (optional)."""
+    """Clear request context at a request or task lifecycle boundary."""
     _request_id.set("")
     _trace_id.set("")
+    _identity.set(_EMPTY_IDENTITY_CONTEXT)

api/core/logging/filters.py

@@ -4,10 +4,7 @@ import contextlib
 import logging
 from typing import override
 
-import flask
-
-from core.logging.context import get_request_id, get_trace_id
-from core.logging.structured_formatter import IdentityDict
+from core.logging.context import get_identity_context, get_request_id, get_trace_id
 
 
 class TraceContextFilter(logging.Filter):
@@ -51,49 +48,16 @@ class TraceContextFilter(logging.Filter):
 
 
 class IdentityContextFilter(logging.Filter):
-    """
-    Filter that adds user identity context to log records.
-    Extracts tenant_id, user_id, and user_type from Flask-Login current_user.
+    """Add an identity snapshot without invoking authentication or database work.
+
+    Logging can run while other libraries hold internal locks, so this filter must
+    only read primitive ContextVar values populated by authentication boundaries.
     """
 
     @override
     def filter(self, record: logging.LogRecord) -> bool:
-        identity = self._extract_identity()
-        record.tenant_id = identity.get("tenant_id", "")
-        record.user_id = identity.get("user_id", "")
-        record.user_type = identity.get("user_type", "")
+        identity = get_identity_context()
+        record.tenant_id = identity.tenant_id
+        record.user_id = identity.user_id
+        record.user_type = identity.user_type
         return True
-
-    def _extract_identity(self) -> IdentityDict:
-        """Extract identity from current_user if in request context."""
-        try:
-            if not flask.has_request_context():
-                return {}
-            from flask_login import current_user
-
-            # Check if user is authenticated using the proxy
-            if not current_user.is_authenticated:
-                return {}
-
-            # Access the underlying user object
-            user = current_user
-
-            from models import Account
-            from models.model import EndUser
-
-            identity: IdentityDict = {}
-
-            match user:
-                case Account():
-                    if user.current_tenant_id:
-                        identity["tenant_id"] = user.current_tenant_id
-                    identity["user_id"] = user.id
-                    identity["user_type"] = "account"
-                case EndUser():
-                    identity["tenant_id"] = user.tenant_id
-                    identity["user_id"] = user.id
-                    identity["user_type"] = user.type or "end_user"
-
-            return identity
-        except Exception:
-            return {}

api/core/rag/docstore/dataset_docstore.py

@@ -6,10 +6,7 @@ from typing import Any
 from sqlalchemy import delete, func, select
 from sqlalchemy.orm import Session
 
-from core.model_manager import ModelManager
-from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from core.rag.models.document import AttachmentDocument, Document
-from graphon.model_runtime.entities.model_entities import ModelType
 from models.dataset import ChildChunk, Dataset, DocumentSegment, SegmentAttachmentBinding
 from models.enums import SegmentType
 
@@ -69,34 +66,22 @@ class DatasetDocumentStore:
 
     def add_documents(
         self,
-        docs: Sequence[Document],
         session: Session,
+        docs: Sequence[Document],
+        token_counts: list[int],
         allow_update: bool = True,
         save_child: bool = False,
-    ):
+    ) -> None:
+        document_token_pairs = list(zip(docs, token_counts, strict=True))
+
         max_position = session.scalar(
             select(func.max(DocumentSegment.position)).where(DocumentSegment.document_id == self._document_id)
         )
 
         if max_position is None:
             max_position = 0
-        embedding_model = None
-        if self._dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
-            model_manager = ModelManager.for_tenant(tenant_id=self._dataset.tenant_id)
-            embedding_model = model_manager.get_model_instance(
-                tenant_id=self._dataset.tenant_id,
-                provider=self._dataset.embedding_model_provider,
-                model_type=ModelType.TEXT_EMBEDDING,
-                model=self._dataset.embedding_model,
-            )
 
-        if embedding_model:
-            page_content_list = [doc.page_content for doc in docs]
-            tokens_list = embedding_model.get_text_embedding_num_tokens(page_content_list)
-        else:
-            tokens_list = [0] * len(docs)
-
-        for doc, tokens in zip(docs, tokens_list):
+        for doc, tokens in document_token_pairs:
             if not isinstance(doc, Document):
                 raise ValueError("doc must be a Document")
 

api/core/rag/embedding/token_counter.py

@@ -0,0 +1,25 @@
+"""Token counting for document segments."""
+
+from core.model_manager import ModelManager
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
+from core.rag.models.document import Document
+from graphon.model_runtime.entities.model_entities import ModelType
+from models.dataset import Dataset
+
+
+def calculate_segment_token_counts(dataset: Dataset, documents: list[Document]) -> list[int]:
+    """Return one token count per document, invoking the embedding model only for high-quality indexes."""
+    if not documents:
+        return []
+
+    if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY:
+        return [0] * len(documents)
+
+    model_manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id)
+    embedding_model = model_manager.get_model_instance(
+        tenant_id=dataset.tenant_id,
+        provider=dataset.embedding_model_provider,
+        model_type=ModelType.TEXT_EMBEDDING,
+        model=dataset.embedding_model,
+    )
+    return embedding_model.get_text_embedding_num_tokens([document.page_content for document in documents])

api/core/rag/index_processor/processor/paragraph_index_processor.py

@@ -19,6 +19,7 @@ from core.rag.cleaner.clean_processor import CleanProcessor
 from core.rag.datasource.keyword.keyword_factory import Keyword
 from core.rag.datasource.vdb.vector_factory import Vector
 from core.rag.docstore.dataset_docstore import DatasetDocumentStore
+from core.rag.embedding.token_counter import calculate_segment_token_counts
 from core.rag.entities import Rule
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.extractor.extract_processor import ExtractProcessor
@@ -244,10 +245,16 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
                         all_multimodal_documents.extend(doc.attachments)
                 documents.append(doc)
         if documents:
+            token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
             # save node to document segment
             doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
             # add document segments
-            doc_store.add_documents(docs=documents, save_child=False, session=session)
+            doc_store.add_documents(
+                session=session,
+                docs=documents,
+                token_counts=token_counts,
+                save_child=False,
+            )
             session.commit()
             if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
                 vector = Vector(dataset, session=session)

api/core/rag/index_processor/processor/parent_child_index_processor.py

@@ -15,6 +15,7 @@ from core.model_manager import ModelInstance
 from core.rag.cleaner.clean_processor import CleanProcessor
 from core.rag.datasource.vdb.vector_factory import Vector
 from core.rag.docstore.dataset_docstore import DatasetDocumentStore
+from core.rag.embedding.token_counter import calculate_segment_token_counts
 from core.rag.entities import ParentMode, Rule
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.extractor.extract_processor import ExtractProcessor
@@ -304,6 +305,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
                 doc.attachments = self._get_content_files(doc, current_user=account, session=session)
             documents.append(doc)
         if documents:
+            token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
             # update document parent mode
             dataset_process_rule = DatasetProcessRule(
                 dataset_id=dataset.id,
@@ -321,7 +323,12 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
             # save node to document segment
             doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
             # add document segments
-            doc_store.add_documents(docs=documents, save_child=True, session=session)
+            doc_store.add_documents(
+                session=session,
+                docs=documents,
+                token_counts=token_counts,
+                save_child=True,
+            )
             session.commit()
             if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
                 all_child_documents = []

api/core/rag/index_processor/processor/qa_index_processor.py

@@ -17,6 +17,7 @@ from core.llm_generator.llm_generator import LLMGenerator
 from core.rag.cleaner.clean_processor import CleanProcessor
 from core.rag.datasource.vdb.vector_factory import Vector
 from core.rag.docstore.dataset_docstore import DatasetDocumentStore
+from core.rag.embedding.token_counter import calculate_segment_token_counts
 from core.rag.entities import Rule
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.extractor.extract_processor import ExtractProcessor
@@ -205,9 +206,15 @@ class QAIndexProcessor(BaseIndexProcessor):
             doc = Document(page_content=qa_chunk.question, metadata=metadata)
             documents.append(doc)
         if documents:
+            token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents)
             # save node to document segment
             doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
-            doc_store.add_documents(docs=documents, save_child=False, session=session)
+            doc_store.add_documents(
+                session=session,
+                docs=documents,
+                token_counts=token_counts,
+                save_child=False,
+            )
             session.commit()
             if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
                 vector = Vector(dataset, session=session)

api/core/tools/__base/tool.py

@@ -17,6 +17,7 @@ from core.tools.entities.tool_entities import (
     ToolParameter,
     ToolProviderType,
 )
+from core.tools.entities.ui_entities import ToolUIMessage
 
 
 class Tool(ABC):
@@ -295,6 +296,13 @@ class Tool(ABC):
             message=ToolInvokeMessage.JsonMessage(json_object=object, suppress_output=suppress_output),
         )
 
+    def create_ui_message(self, ui_message: ToolUIMessage | dict[str, Any]) -> ToolInvokeMessage:
+        """Create a validated, model-invisible UI message for chat clients."""
+        return ToolInvokeMessage(
+            type=ToolInvokeMessage.MessageType.UI,
+            message=ToolUIMessage.model_validate(ui_message),
+        )
+
     def create_variable_message(
         self, variable_name: str, variable_value: Any, stream: bool = False
     ) -> ToolInvokeMessage:

api/core/tools/builtin_tool/providers/time/tools/current_time.py

@@ -7,6 +7,9 @@ from sqlalchemy.orm import Session
 
 from core.tools.builtin_tool.tool import BuiltinTool
 from core.tools.entities.tool_entities import ToolInvokeMessage
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID, A2UI_PROTOCOL_VERSION
+
+_CURRENT_TIME_SURFACE_ID = "current-time"
 
 
 class CurrentTimeTool(BuiltinTool):
@@ -20,19 +23,60 @@ class CurrentTimeTool(BuiltinTool):
         app_id: str | None = None,
         message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
-        """
-        invoke tools
-        """
-        # get timezone
-        tz = tool_parameters.get("timezone", "UTC")
-        fm = tool_parameters.get("format") or "%Y-%m-%d %H:%M:%S %Z"
-        if tz == "UTC":
-            yield self.create_text_message(f"{datetime.now(UTC).strftime(fm)}")
-            return
+        """Return the formatted time for the model and a standard time card for chat clients.
 
-        try:
-            tz = pytz_timezone(tz)
-        except Exception:
-            yield self.create_text_message(f"Invalid timezone: {tz}")
-            return
-        yield self.create_text_message(f"{datetime.now(tz).strftime(fm)}")
+        Invalid timezone values retain the existing text-only error response.
+        """
+        timezone_name = tool_parameters.get("timezone", "UTC")
+        fm = tool_parameters.get("format") or "%Y-%m-%d %H:%M:%S %Z"
+        if timezone_name == "UTC":
+            current_time = datetime.now(UTC)
+        else:
+            try:
+                timezone_info = pytz_timezone(timezone_name)
+            except Exception:
+                yield self.create_text_message(f"Invalid timezone: {timezone_name}")
+                return
+            current_time = datetime.now(timezone_info)
+
+        formatted_time = current_time.strftime(fm)
+        yield self.create_text_message(formatted_time)
+        yield self.create_ui_message(
+            {
+                "messages": [
+                    {
+                        "version": A2UI_PROTOCOL_VERSION,
+                        "createSurface": {
+                            "surfaceId": _CURRENT_TIME_SURFACE_ID,
+                            "catalogId": A2UI_CATALOG_ID,
+                        },
+                    },
+                    {
+                        "version": A2UI_PROTOCOL_VERSION,
+                        "updateDataModel": {
+                            "surfaceId": _CURRENT_TIME_SURFACE_ID,
+                            "value": {"currentTime": current_time.isoformat()},
+                        },
+                    },
+                    {
+                        "version": A2UI_PROTOCOL_VERSION,
+                        "updateComponents": {
+                            "surfaceId": _CURRENT_TIME_SURFACE_ID,
+                            "components": [
+                                {
+                                    "id": "root",
+                                    "component": "Card",
+                                    "children": ["time"],
+                                },
+                                {
+                                    "id": "time",
+                                    "component": "DateTime",
+                                    "value": {"path": "/currentTime"},
+                                    "format": "datetime",
+                                },
+                            ],
+                        },
+                    },
+                ],
+            }
+        )

api/core/tools/entities/tool_entities.py

@@ -32,6 +32,7 @@ from core.plugin.entities.parameters import (
 from core.rag.entities import RetrievalSourceMetadata
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.constants import TOOL_SELECTOR_MODEL_IDENTITY
+from core.tools.entities.ui_entities import ToolUIMessage
 
 
 class EmojiIconDict(TypedDict):
@@ -240,13 +241,15 @@ class ToolInvokeMessage(BaseModel):
         LOG = auto()
         BLOB_CHUNK = auto()
         RETRIEVER_RESOURCES = auto()
+        UI = auto()
 
     type: MessageType = MessageType.TEXT
     """
         plain text, image url or link url
     """
     message: (
-        JsonMessage
+        ToolUIMessage
+        | JsonMessage
         | TextMessage
         | BlobChunkMessage
         | BlobMessage
@@ -275,6 +278,8 @@ class ToolInvokeMessage(BaseModel):
                     v = {"json_object": v}
             elif msg_type == cls.MessageType.FILE:
                 v = {"file_marker": "file_marker"}
+            elif msg_type == cls.MessageType.UI:
+                v = ToolUIMessage.model_validate(v)
 
         return v
 

api/core/tools/entities/ui_entities.py

@@ -0,0 +1,675 @@
+"""Validated, read-only UI messages emitted by tools.
+
+The public wire format is the A2UI v0.9.1 flat message shape, narrowed to a
+Dify-owned catalog. Tool authors can bind display values to the surface data
+model, but cannot supply executable actions, HTML, styles, themes, or custom
+components. A complete ``ToolUIMessage`` describes exactly one surface and is
+validated as a bounded, self-contained component graph before it crosses into
+chat streaming. Sequential data-model patches are materialized with the same
+object/array upsert semantics as the web renderer so cumulative limits cannot
+be bypassed with individually small updates.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import uuid
+from collections.abc import Sequence
+from enum import StrEnum
+from typing import Annotated, Literal
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    JsonValue,
+    SerializerFunctionWrapHandler,
+    field_validator,
+    model_serializer,
+    model_validator,
+)
+
+A2UI_PROTOCOL = "a2ui"
+A2UI_PROTOCOL_VERSION = "v0.9.1"
+A2UI_CATALOG_ID = "https://dify.ai/a2ui/catalog/v1"
+DIFY_UI_JSON_ENVELOPE_KEY = "__dify_ui__"
+
+MAX_UI_MESSAGES = 64
+MAX_UI_COMPONENTS = 100
+MAX_UI_STRING_LENGTH = 4096
+MAX_UI_PAYLOAD_BYTES = 128 * 1024
+MAX_DATA_MODEL_DEPTH = 16
+MAX_DATA_MODEL_NODES = 2000
+MAX_DATA_MODEL_ARRAY_INDEX = 1000
+MAX_JSON_POINTER_SEGMENTS = 16
+MAX_UI_PARTS_PER_MESSAGE = 16
+MAX_UI_PARTS_PAYLOAD_BYTES = 512 * 1024
+
+_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}$")
+_ARRAY_INDEX_PATTERN = re.compile(r"^(?:0|[1-9]\d*)$")
+_DANGEROUS_POINTER_SEGMENTS = {"__proto__", "constructor", "prototype"}
+_DANGEROUS_DATA_KEYS = _DANGEROUS_POINTER_SEGMENTS
+_MAX_HISTORY_UI_PART_CANDIDATES = MAX_UI_PARTS_PER_MESSAGE * 4
+_UI_PART_ID_NAMESPACE = uuid.uuid5(uuid.NAMESPACE_URL, "https://dify.ai/a2ui/part-id/v1")
+
+
+class _StrictModel(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        populate_by_name=True,
+        serialize_by_alias=True,
+        allow_inf_nan=False,
+    )
+
+
+class A2UIDataBinding(_StrictModel):
+    """A JSON Pointer into the surface data model."""
+
+    path: str
+
+    @field_validator("path")
+    @classmethod
+    def _validate_path(cls, value: str) -> str:
+        return _validate_json_pointer(value, label="data binding path")
+
+
+type DynamicString = str | A2UIDataBinding
+type DynamicNumber = int | float | A2UIDataBinding
+type DynamicScalar = str | int | float | bool | None | A2UIDataBinding
+
+
+class A2UIComponentType(StrEnum):
+    CARD = "Card"
+    ROW = "Row"
+    COLUMN = "Column"
+    TEXT = "Text"
+    ICON = "Icon"
+    DIVIDER = "Divider"
+    BADGE = "Badge"
+    METRIC = "Metric"
+    DATE_TIME = "DateTime"
+    PROGRESS = "Progress"
+    KEY_VALUE = "KeyValue"
+
+
+class A2UIComponent(_StrictModel):
+    """One flat component entry from the Dify catalog."""
+
+    id: str
+    component: A2UIComponentType
+    children: list[str] | None = None
+    title: DynamicString | None = None
+    gap: Literal["small", "medium", "large"] | None = None
+    align: Literal["start", "center", "end"] | None = None
+    text: DynamicString | None = None
+    variant: Literal["body", "caption"] | None = None
+    name: (
+        Literal[
+            "clock",
+            "cloud",
+            "sun",
+            "rain",
+            "snow",
+            "wind",
+            "thermometer",
+            "calendar",
+            "location",
+        ]
+        | None
+    ) = None
+    tone: Literal["neutral", "info", "success", "warning", "critical"] | None = None
+    label: DynamicString | None = None
+    value: DynamicScalar | None = None
+    unit: DynamicString | None = None
+    format: Literal["date", "time", "datetime"] | None = None
+    max: DynamicNumber | None = None
+
+    @field_validator("id")
+    @classmethod
+    def _validate_id(cls, value: str) -> str:
+        if not _ID_PATTERN.fullmatch(value):
+            raise ValueError("component id contains unsupported characters or is too long")
+        return value
+
+    @field_validator("children")
+    @classmethod
+    def _validate_children(cls, value: list[str] | None) -> list[str] | None:
+        if value is None:
+            return None
+        if len(value) > MAX_UI_COMPONENTS:
+            raise ValueError("component has too many children")
+        if len(value) != len(set(value)):
+            raise ValueError("component children must not contain duplicate ids")
+        for child_id in value:
+            if not _ID_PATTERN.fullmatch(child_id):
+                raise ValueError("child component id contains unsupported characters or is too long")
+        return value
+
+    @model_validator(mode="after")
+    def _validate_component_props(self) -> A2UIComponent:
+        present = self.model_fields_set - {"id", "component"}
+        allowed: dict[A2UIComponentType, set[str]] = {
+            A2UIComponentType.CARD: {"children", "title"},
+            A2UIComponentType.ROW: {"children", "gap", "align"},
+            A2UIComponentType.COLUMN: {"children", "gap"},
+            A2UIComponentType.TEXT: {"text", "variant"},
+            A2UIComponentType.ICON: {"name"},
+            A2UIComponentType.DIVIDER: set(),
+            A2UIComponentType.BADGE: {"text", "tone"},
+            A2UIComponentType.METRIC: {"label", "value", "unit"},
+            A2UIComponentType.DATE_TIME: {"value", "format"},
+            A2UIComponentType.PROGRESS: {"value", "max", "label"},
+            A2UIComponentType.KEY_VALUE: {"label", "value"},
+        }
+        required: dict[A2UIComponentType, set[str]] = {
+            A2UIComponentType.CARD: {"children"},
+            A2UIComponentType.ROW: {"children"},
+            A2UIComponentType.COLUMN: {"children"},
+            A2UIComponentType.TEXT: {"text"},
+            A2UIComponentType.ICON: {"name"},
+            A2UIComponentType.DIVIDER: set(),
+            A2UIComponentType.BADGE: {"text"},
+            A2UIComponentType.METRIC: {"label", "value"},
+            A2UIComponentType.DATE_TIME: {"value"},
+            A2UIComponentType.PROGRESS: {"value", "label"},
+            A2UIComponentType.KEY_VALUE: {"label", "value"},
+        }
+        unexpected = present - allowed[self.component]
+        missing = required[self.component] - present
+        if unexpected:
+            raise ValueError(f"{self.component} does not support properties: {sorted(unexpected)}")
+        if missing:
+            raise ValueError(f"{self.component} requires properties: {sorted(missing)}")
+        nullable_props = (
+            {"value"} if self.component in {A2UIComponentType.METRIC, A2UIComponentType.KEY_VALUE} else set()
+        )
+        null_props = {prop for prop in present if getattr(self, prop) is None and prop not in nullable_props}
+        if null_props:
+            raise ValueError(f"{self.component} properties cannot be null: {sorted(null_props)}")
+        if self.component == A2UIComponentType.DATE_TIME and not isinstance(self.value, str | A2UIDataBinding):
+            raise ValueError("DateTime value must be a string or data binding")
+        if self.component == A2UIComponentType.PROGRESS and (
+            isinstance(self.value, bool) or not isinstance(self.value, int | float | A2UIDataBinding)
+        ):
+            raise ValueError("Progress value must be a number or data binding")
+        if (
+            self.component == A2UIComponentType.PROGRESS
+            and "max" in present
+            and (isinstance(self.max, bool) or not isinstance(self.max, int | float | A2UIDataBinding))
+        ):
+            raise ValueError("Progress max must be a number or data binding")
+        return self
+
+    @model_serializer(mode="wrap")
+    def _serialize_component(self, handler: SerializerFunctionWrapHandler) -> dict[str, JsonValue]:
+        serialized = handler(self)
+        return {
+            key: value
+            for key, value in serialized.items()
+            if key in {"id", "component"} or key in self.model_fields_set
+        }
+
+
+class A2UICreateSurface(_StrictModel):
+    surface_id: str = Field(alias="surfaceId")
+    catalog_id: Literal["https://dify.ai/a2ui/catalog/v1"] = Field(alias="catalogId")
+
+    @field_validator("surface_id")
+    @classmethod
+    def _validate_surface_id(cls, value: str) -> str:
+        if not _ID_PATTERN.fullmatch(value):
+            raise ValueError("surface id contains unsupported characters or is too long")
+        return value
+
+
+class A2UIUpdateComponents(_StrictModel):
+    surface_id: str = Field(alias="surfaceId")
+    components: Annotated[list[A2UIComponent], Field(min_length=1, max_length=MAX_UI_COMPONENTS)]
+
+    @field_validator("surface_id")
+    @classmethod
+    def _validate_surface_id(cls, value: str) -> str:
+        if not _ID_PATTERN.fullmatch(value):
+            raise ValueError("surface id contains unsupported characters or is too long")
+        return value
+
+    @field_validator("components")
+    @classmethod
+    def _validate_unique_component_ids(cls, value: list[A2UIComponent]) -> list[A2UIComponent]:
+        ids = [component.id for component in value]
+        if len(ids) != len(set(ids)):
+            raise ValueError("updateComponents contains duplicate component ids")
+        return value
+
+
+class A2UIUpdateDataModel(_StrictModel):
+    surface_id: str = Field(alias="surfaceId")
+    value: JsonValue
+    path: str | None = None
+
+    @field_validator("surface_id")
+    @classmethod
+    def _validate_surface_id(cls, value: str) -> str:
+        if not _ID_PATTERN.fullmatch(value):
+            raise ValueError("surface id contains unsupported characters or is too long")
+        return value
+
+    @field_validator("value")
+    @classmethod
+    def _validate_value(cls, value: JsonValue) -> JsonValue:
+        _validate_data_model_value(value)
+        return value
+
+    @field_validator("path")
+    @classmethod
+    def _validate_path(cls, value: str | None) -> str | None:
+        if value is None:
+            return value
+        return _validate_json_pointer(value, label="data model path", enforce_array_index_limit=True)
+
+    @model_serializer(mode="wrap")
+    def _serialize_update(self, handler: SerializerFunctionWrapHandler) -> dict[str, JsonValue]:
+        serialized = handler(self)
+        if "path" not in self.model_fields_set:
+            serialized.pop("path", None)
+        return serialized
+
+
+class A2UIDeleteSurface(_StrictModel):
+    surface_id: str = Field(alias="surfaceId")
+
+    @field_validator("surface_id")
+    @classmethod
+    def _validate_surface_id(cls, value: str) -> str:
+        if not _ID_PATTERN.fullmatch(value):
+            raise ValueError("surface id contains unsupported characters or is too long")
+        return value
+
+
+class A2UIMessage(_StrictModel):
+    """A single A2UI v0.9.1 server message."""
+
+    version: Literal["v0.9.1"]
+    create_surface: A2UICreateSurface | None = Field(default=None, alias="createSurface")
+    update_components: A2UIUpdateComponents | None = Field(default=None, alias="updateComponents")
+    update_data_model: A2UIUpdateDataModel | None = Field(default=None, alias="updateDataModel")
+    delete_surface: A2UIDeleteSurface | None = Field(default=None, alias="deleteSurface")
+
+    @model_validator(mode="after")
+    def _validate_exactly_one_operation(self) -> A2UIMessage:
+        operations = (
+            self.create_surface,
+            self.update_components,
+            self.update_data_model,
+            self.delete_surface,
+        )
+        if sum(operation is not None for operation in operations) != 1:
+            raise ValueError("A2UI message must contain exactly one operation")
+        return self
+
+    @model_serializer(mode="wrap")
+    def _serialize_message(self, handler: SerializerFunctionWrapHandler) -> dict[str, JsonValue]:
+        return {key: value for key, value in handler(self).items() if value is not None}
+
+    @property
+    def surface_id(self) -> str:
+        operation = self.create_surface or self.update_components or self.update_data_model or self.delete_surface
+        assert operation is not None
+        return operation.surface_id
+
+
+class ToolUIMessage(_StrictModel):
+    """A complete, validated UI surface emitted by one tool result.
+
+    Validation covers both the component graph and every materialized
+    data-model revision, because clients render the message sequence in order.
+    """
+
+    protocol: Literal["a2ui"] = "a2ui"
+    protocol_version: Literal["v0.9.1"] = "v0.9.1"
+    messages: Annotated[list[A2UIMessage], Field(min_length=1, max_length=MAX_UI_MESSAGES)]
+    fallback: str | None = None
+
+    @field_validator("fallback")
+    @classmethod
+    def _validate_fallback(cls, value: str | None) -> str | None:
+        if value is not None and len(value) > MAX_UI_STRING_LENGTH:
+            raise ValueError("UI fallback is too long")
+        return value
+
+    @model_validator(mode="after")
+    def _validate_surface(self) -> ToolUIMessage:
+        first = self.messages[0]
+        if first.create_surface is None:
+            raise ValueError("first A2UI message must be createSurface")
+        if any(message.create_surface is not None for message in self.messages[1:]):
+            raise ValueError("createSurface may only appear as the first message")
+        delete_indexes = [index for index, message in enumerate(self.messages) if message.delete_surface]
+        if delete_indexes and delete_indexes != [len(self.messages) - 1]:
+            raise ValueError("deleteSurface may only appear once as the final message")
+
+        surface_id = first.surface_id
+        if any(message.surface_id != surface_id for message in self.messages):
+            raise ValueError("all A2UI messages in a tool UI message must target the same surface")
+
+        components: dict[str, A2UIComponent] = {}
+        data_model: JsonValue = {}
+        root_seen_before_delete = False
+        for message in self.messages[1:]:
+            if message.update_components is not None:
+                for component in message.update_components.components:
+                    components[component.id] = component
+                if len(components) > MAX_UI_COMPONENTS:
+                    raise ValueError("UI surface has too many components")
+            if message.update_data_model is not None:
+                data_model = _apply_data_model_update(data_model, message.update_data_model)
+                _validate_data_model_value(data_model)
+            if message.delete_surface is not None:
+                root_seen_before_delete = "root" in components
+
+        if "root" not in components and not root_seen_before_delete:
+            raise ValueError("UI surface must define a component with id 'root'")
+
+        _validate_component_graph(components)
+        payload = self.model_dump(mode="json", by_alias=True, exclude_none=True)
+        _validate_bounded_json(payload)
+        return self
+
+    @model_serializer(mode="wrap")
+    def _serialize_ui_message(self, handler: SerializerFunctionWrapHandler) -> dict[str, JsonValue]:
+        return {key: value for key, value in handler(self).items() if value is not None}
+
+    @property
+    def surface_id(self) -> str:
+        return self.messages[0].surface_id
+
+
+class MessageUIPart(_StrictModel):
+    """SSE/history representation of one tool-owned UI surface revision."""
+
+    part_id: Annotated[str, Field(min_length=1, max_length=512)]
+    sequence: Annotated[int, Field(ge=1)]
+    protocol: Literal["a2ui"] = "a2ui"
+    protocol_version: Literal["v0.9.1"] = "v0.9.1"
+    messages: Annotated[list[A2UIMessage], Field(min_length=1, max_length=MAX_UI_MESSAGES)]
+    fallback: str | None = None
+
+    @model_validator(mode="after")
+    def _validate_messages(self) -> MessageUIPart:
+        ToolUIMessage(
+            protocol=self.protocol,
+            protocol_version=self.protocol_version,
+            messages=self.messages,
+            fallback=self.fallback,
+        )
+        return self
+
+    @model_serializer(mode="wrap")
+    def _serialize_part(self, handler: SerializerFunctionWrapHandler) -> dict[str, JsonValue]:
+        return {key: value for key, value in handler(self).items() if value is not None}
+
+    @classmethod
+    def from_tool_ui_message(
+        cls,
+        *,
+        part_id: str,
+        sequence: int,
+        ui_message: ToolUIMessage,
+    ) -> MessageUIPart:
+        return cls(
+            part_id=part_id,
+            sequence=sequence,
+            protocol=ui_message.protocol,
+            protocol_version=ui_message.protocol_version,
+            messages=ui_message.messages,
+            fallback=ui_message.fallback,
+        )
+
+
+def extract_ui_message_from_json(value: JsonValue) -> ToolUIMessage | None:
+    """Recognize the reserved compatibility envelope used by older SDKs."""
+
+    if not isinstance(value, dict) or set(value) != {DIFY_UI_JSON_ENVELOPE_KEY}:
+        return None
+    return ToolUIMessage.model_validate(value[DIFY_UI_JSON_ENVELOPE_KEY])
+
+
+def parse_tool_ui_messages(value: object) -> list[ToolUIMessage]:
+    """Validate a list received through the Agent backend metadata channel."""
+
+    if not isinstance(value, list):
+        raise ValueError("dify_ui_messages metadata must be a list")
+    if len(value) > MAX_UI_PARTS_PER_MESSAGE:
+        raise ValueError(f"tool UI batch cannot contain more than {MAX_UI_PARTS_PER_MESSAGE} messages")
+    messages = [ToolUIMessage.model_validate(item) for item in value]
+    validate_tool_ui_message_batch(messages)
+    return messages
+
+
+def build_ui_part_id(namespace: str, surface_id: str) -> str:
+    """Build a stable, bounded ID from an unambiguous tool/surface tuple."""
+    if not isinstance(namespace, str) or not isinstance(surface_id, str):
+        raise TypeError("UI part namespace and surface id must be strings")
+    identity = json.dumps((namespace, surface_id), ensure_ascii=True, separators=(",", ":"))
+    return f"ui-{uuid.uuid5(_UI_PART_ID_NAMESPACE, identity)}"
+
+
+def validate_tool_ui_message_batch(messages: Sequence[ToolUIMessage]) -> None:
+    """Enforce the per-tool-call UI batch budget before queue publication."""
+    if len(messages) > MAX_UI_PARTS_PER_MESSAGE:
+        raise ValueError(f"tool UI batch cannot contain more than {MAX_UI_PARTS_PER_MESSAGE} messages")
+    if _serialized_model_list_size(messages) > MAX_UI_PARTS_PAYLOAD_BYTES:
+        raise ValueError(f"tool UI batch payload cannot exceed {MAX_UI_PARTS_PAYLOAD_BYTES} bytes")
+
+
+def validate_ui_part_batch(parts: Sequence[MessageUIPart]) -> None:
+    """Enforce persisted/current assistant-message UI budgets."""
+    if len(parts) > MAX_UI_PARTS_PER_MESSAGE:
+        raise ValueError(f"assistant message cannot contain more than {MAX_UI_PARTS_PER_MESSAGE} UI parts")
+    part_ids = [part.part_id for part in parts]
+    if len(part_ids) != len(set(part_ids)):
+        raise ValueError("assistant message UI parts must have distinct part ids")
+    if _serialized_model_list_size(parts) > MAX_UI_PARTS_PAYLOAD_BYTES:
+        raise ValueError(f"assistant message UI parts cannot exceed {MAX_UI_PARTS_PAYLOAD_BYTES} bytes")
+
+
+def upsert_ui_part(parts: Sequence[MessageUIPart], part: MessageUIPart) -> list[MessageUIPart] | None:
+    """Return a bounded new part list, or ``None`` for a stale revision."""
+    candidate = list(parts)
+    for index, existing_part in enumerate(candidate):
+        if existing_part.part_id != part.part_id:
+            continue
+        if existing_part.sequence >= part.sequence:
+            return None
+        candidate[index] = part
+        break
+    else:
+        candidate.append(part)
+
+    validate_ui_part_batch(candidate)
+    return candidate
+
+
+def parse_history_ui_parts(value: object) -> list[MessageUIPart]:
+    """Best-effort recovery of bounded UI parts from untrusted historical metadata."""
+    if not isinstance(value, list):
+        return []
+
+    parts: list[MessageUIPart] = []
+    for item in value[:_MAX_HISTORY_UI_PART_CANDIDATES]:
+        try:
+            part = MessageUIPart.model_validate(item)
+            updated_parts = upsert_ui_part(parts, part)
+        except (TypeError, ValueError):
+            continue
+        if updated_parts is not None:
+            parts = updated_parts
+    return parts
+
+
+def _validate_component_graph(components: dict[str, A2UIComponent]) -> None:
+    for component in components.values():
+        for child_id in component.children or []:
+            if child_id not in components:
+                raise ValueError(f"component {component.id!r} references unknown child {child_id!r}")
+
+    visiting: set[str] = set()
+    visited: set[str] = set()
+
+    def visit(component_id: str) -> None:
+        if component_id in visiting:
+            raise ValueError("UI component graph contains a cycle")
+        if component_id in visited:
+            return
+        visiting.add(component_id)
+        for child_id in components[component_id].children or []:
+            visit(child_id)
+        visiting.remove(component_id)
+        visited.add(component_id)
+
+    for component_id in components:
+        visit(component_id)
+
+    parent_by_child: dict[str, str] = {}
+    for component in components.values():
+        for child_id in component.children or []:
+            if child_id == "root":
+                raise ValueError("root component cannot be referenced as a child")
+            existing_parent = parent_by_child.get(child_id)
+            if existing_parent is not None:
+                raise ValueError(
+                    f"component {child_id!r} has multiple parents: {existing_parent!r} and {component.id!r}"
+                )
+            parent_by_child[child_id] = component.id
+
+    reachable = {"root"}
+    pending = ["root"]
+    while pending:
+        component_id = pending.pop()
+        for child_id in components[component_id].children or []:
+            if child_id in reachable:
+                continue
+            reachable.add(child_id)
+            pending.append(child_id)
+
+    unreachable = set(components) - reachable
+    if unreachable:
+        raise ValueError(f"UI component graph contains components unreachable from root: {sorted(unreachable)}")
+
+
+def _serialized_model_list_size(values: Sequence[BaseModel]) -> int:
+    payload = [value.model_dump(mode="json", by_alias=True, exclude_none=True) for value in values]
+    return len(json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode())
+
+
+def _validate_json_pointer(
+    value: str,
+    *,
+    label: str,
+    enforce_array_index_limit: bool = False,
+) -> str:
+    if not value.startswith("/"):
+        raise ValueError(f"{label} must be a JSON Pointer beginning with '/'")
+    if len(value) > 256:
+        raise ValueError(f"{label} is too long")
+    if re.search(r"~(?![01])", value):
+        raise ValueError(f"{label} contains an invalid JSON Pointer escape")
+    decoded_segments = _decode_json_pointer(value)
+    if len(decoded_segments) > MAX_JSON_POINTER_SEGMENTS:
+        raise ValueError(f"{label} contains too many segments")
+    if set(decoded_segments).intersection(_DANGEROUS_POINTER_SEGMENTS):
+        raise ValueError(f"{label} contains a forbidden segment")
+    if enforce_array_index_limit and any(
+        segment.isdecimal() and int(segment) > MAX_DATA_MODEL_ARRAY_INDEX for segment in decoded_segments
+    ):
+        raise ValueError(f"{label} contains an array index larger than {MAX_DATA_MODEL_ARRAY_INDEX}")
+    return value
+
+
+def _decode_json_pointer(value: str) -> list[str]:
+    if value == "/":
+        return []
+    return [segment.replace("~1", "/").replace("~0", "~") for segment in value.split("/")[1:]]
+
+
+def _apply_data_model_update(current: JsonValue, update: A2UIUpdateDataModel) -> JsonValue:
+    path = update.path
+    if path is None or path == "/":
+        return update.value
+
+    segments = _decode_json_pointer(path)
+
+    def apply_at(node: JsonValue | None, segment_index: int) -> JsonValue:
+        if segment_index == len(segments):
+            return update.value
+
+        segment = segments[segment_index]
+        is_array_segment = _ARRAY_INDEX_PATTERN.fullmatch(segment) is not None
+        if isinstance(node, list) or (not isinstance(node, dict) and is_array_segment):
+            if not is_array_segment:
+                raise ValueError("data model array paths must use canonical non-negative indexes")
+            index = int(segment)
+            if index > MAX_DATA_MODEL_ARRAY_INDEX:
+                raise ValueError(f"data model array index cannot exceed {MAX_DATA_MODEL_ARRAY_INDEX}")
+
+            next_node = list(node) if isinstance(node, list) else []
+            if index > len(next_node):
+                raise ValueError("data model array index cannot create a gap")
+            child = next_node[index] if index < len(next_node) else None
+            updated_child = apply_at(child, segment_index + 1)
+            if index == len(next_node):
+                next_node.append(updated_child)
+            else:
+                next_node[index] = updated_child
+            return next_node
+
+        next_node = dict(node) if isinstance(node, dict) else {}
+        next_node[segment] = apply_at(next_node.get(segment), segment_index + 1)
+        return next_node
+
+    return apply_at(current, 0)
+
+
+def _validate_data_model_value(value: JsonValue) -> None:
+    nodes = 0
+
+    def walk(node: JsonValue, *, depth: int) -> None:
+        nonlocal nodes
+        nodes += 1
+        if nodes > MAX_DATA_MODEL_NODES:
+            raise ValueError("data model value has too many nodes")
+        if depth > MAX_DATA_MODEL_DEPTH:
+            raise ValueError("data model value is nested too deeply")
+        if isinstance(node, list):
+            for item in node:
+                walk(item, depth=depth + 1)
+            return
+        if isinstance(node, dict):
+            unsafe_keys = set(node).intersection(_DANGEROUS_DATA_KEYS)
+            if unsafe_keys:
+                raise ValueError(f"data model value contains forbidden keys: {sorted(unsafe_keys)}")
+            for item in node.values():
+                walk(item, depth=depth + 1)
+
+    walk(value, depth=0)
+
+
+def _validate_bounded_json(value: JsonValue) -> None:
+    def walk(node: JsonValue) -> None:
+        if isinstance(node, str):
+            if len(node) > MAX_UI_STRING_LENGTH:
+                raise ValueError("UI payload contains a string that is too long")
+            return
+        if isinstance(node, list):
+            for item in node:
+                walk(item)
+            return
+        if isinstance(node, dict):
+            for item in node.values():
+                walk(item)
+
+    walk(value)
+    encoded = json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode()
+    if len(encoded) > MAX_UI_PAYLOAD_BYTES:
+        raise ValueError("UI payload is too large")

api/core/tools/tool_engine.py

@@ -3,6 +3,7 @@ import json
 import logging
 from collections.abc import Generator, Iterable
 from copy import deepcopy
+from dataclasses import dataclass
 from datetime import UTC, datetime
 from mimetypes import guess_type
 from typing import Any, Union, cast
@@ -21,6 +22,12 @@ from core.tools.entities.tool_entities import (
     ToolInvokeMeta,
     ToolParameter,
 )
+from core.tools.entities.ui_entities import (
+    DIFY_UI_JSON_ENVELOPE_KEY,
+    ToolUIMessage,
+    extract_ui_message_from_json,
+    validate_tool_ui_message_batch,
+)
 from core.tools.errors import (
     ToolEngineInvokeError,
     ToolInvokeError,
@@ -40,6 +47,16 @@ from models.model import Message, MessageFile
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class ToolAgentInvokeResult:
+    """Agent-facing tool result with UI kept outside the model observation."""
+
+    observation: str
+    message_files: list[str]
+    ui_messages: list[ToolUIMessage]
+    meta: ToolInvokeMeta
+
+
 class ToolEngine:
     """
     Tool runtime engine take care of the tool executions.
@@ -59,9 +76,13 @@ class ToolEngine:
         conversation_id: str | None = None,
         app_id: str | None = None,
         message_id: str | None = None,
-    ) -> tuple[str, list[str], ToolInvokeMeta]:
+    ) -> ToolAgentInvokeResult:
         """
-        Agent invokes the tool with the given arguments.
+        Invoke a tool for an agent.
+
+        UI messages are validated and returned on a dedicated channel. They are
+        intentionally excluded from both the LLM observation and tracing
+        callback output.
         """
         # check if arguments is a string
         if isinstance(tool_parameters, str):
@@ -103,7 +124,7 @@ class ToolEngine:
                 conversation_id=message.conversation_id,
             )
 
-            message_list = list(messages)
+            message_list, ui_messages = ToolEngine.collect_agent_messages(messages)
 
             # extract binary data from tool invoke message
             binary_files = ToolEngine._extract_tool_response_binary_and_text(message_list)
@@ -126,7 +147,12 @@ class ToolEngine:
             )
 
             # transform tool invoke message to get LLM friendly message
-            return plain_text, message_files, meta
+            return ToolAgentInvokeResult(
+                observation=plain_text,
+                message_files=message_files,
+                ui_messages=ui_messages,
+                meta=meta,
+            )
         except ToolProviderCredentialValidationError as e:
             logger.error(e, exc_info=True)
             error_response = "Please check your tool provider credentials"
@@ -148,13 +174,23 @@ class ToolEngine:
             error_response = f"tool invoke error: {meta.error}"
             agent_tool_callback.on_tool_error(e)
             logger.error(e, exc_info=True)
-            return error_response, [], meta
+            return ToolAgentInvokeResult(
+                observation=error_response,
+                message_files=[],
+                ui_messages=[],
+                meta=meta,
+            )
         except Exception as e:
             error_response = f"unknown error: {e}"
             agent_tool_callback.on_tool_error(e)
             logger.error(e, exc_info=True)
 
-        return error_response, [], ToolInvokeMeta.error_instance(error_response)
+        return ToolAgentInvokeResult(
+            observation=error_response,
+            message_files=[],
+            ui_messages=[],
+            meta=ToolInvokeMeta.error_instance(error_response),
+        )
 
     @staticmethod
     def generic_invoke(
@@ -197,7 +233,7 @@ class ToolEngine:
                 tool_outputs=response,
             )
 
-            return response
+            return ToolEngine.normalize_ui_messages(response)
         except Exception as e:
             workflow_tool_callback.on_tool_error(e)
             raise e
@@ -266,7 +302,10 @@ class ToolEngine:
                         ensure_ascii=False,
                     )
                 )
-            elif response.type == ToolInvokeMessage.MessageType.VARIABLE:
+            elif response.type in {
+                ToolInvokeMessage.MessageType.VARIABLE,
+                ToolInvokeMessage.MessageType.UI,
+            }:
                 continue
             else:
                 parts.append(str(response.message))
@@ -278,6 +317,124 @@ class ToolEngine:
 
         return "".join(parts)
 
+    @staticmethod
+    def normalize_ui_messages(messages: Iterable[ToolInvokeMessage]) -> Generator[ToolInvokeMessage, None, None]:
+        """Convert reserved variable/JSON compatibility envelopes into native UI."""
+        for message in messages:
+            if message.type == ToolInvokeMessage.MessageType.VARIABLE:
+                variable_message = cast(ToolInvokeMessage.VariableMessage, message.message)
+                if variable_message.variable_name != DIFY_UI_JSON_ENVELOPE_KEY:
+                    yield message
+                    continue
+                yield ToolInvokeMessage(
+                    type=ToolInvokeMessage.MessageType.UI,
+                    message=ToolUIMessage.model_validate(variable_message.variable_value),
+                    meta=message.meta,
+                )
+                continue
+
+            if message.type != ToolInvokeMessage.MessageType.JSON:
+                yield message
+                continue
+
+            json_message = cast(ToolInvokeMessage.JsonMessage, message.message)
+            ui_message = extract_ui_message_from_json(json_message.json_object)
+            if ui_message is None:
+                yield message
+                continue
+            yield ToolInvokeMessage(type=ToolInvokeMessage.MessageType.UI, message=ui_message, meta=message.meta)
+
+    @staticmethod
+    def collect_agent_messages(
+        messages: Iterable[ToolInvokeMessage],
+    ) -> tuple[list[ToolInvokeMessage], list[ToolUIMessage]]:
+        """Collect a bounded UI batch while preserving all non-UI observations.
+
+        A malformed or over-budget UI message invalidates the complete UI
+        batch. Remaining UI is skipped without parsing, but the input iterable
+        is still consumed so later text and file messages are retained.
+        """
+        normalized_messages: list[ToolInvokeMessage] = []
+        ui_messages: list[ToolUIMessage] = []
+        ui_batch_rejected = False
+
+        for message in messages:
+            is_ui_transport = ToolEngine._is_ui_transport_message(message)
+            if ui_batch_rejected and is_ui_transport:
+                continue
+
+            try:
+                normalized = ToolEngine.normalize_ui_messages([message])
+            except (TypeError, ValueError):
+                if not is_ui_transport:
+                    raise
+                # Reserved UI envelopes must never fall back to JSON/variable
+                # observations, even when malformed.
+                ui_batch_rejected = True
+                ui_messages.clear()
+                normalized_messages = [
+                    normalized_message
+                    for normalized_message in normalized_messages
+                    if normalized_message.type != ToolInvokeMessage.MessageType.UI
+                ]
+                logger.warning("Ignored invalid reserved tool UI message", exc_info=True)
+                continue
+
+            try:
+                for normalized_message in normalized:
+                    if normalized_message.type != ToolInvokeMessage.MessageType.UI:
+                        normalized_messages.append(normalized_message)
+                        continue
+                    if ui_batch_rejected:
+                        continue
+
+                    ui_message = cast(ToolUIMessage, normalized_message.message)
+                    candidate = [*ui_messages, ui_message]
+                    try:
+                        validate_tool_ui_message_batch(candidate)
+                    except ValueError:
+                        ui_batch_rejected = True
+                        ui_messages.clear()
+                        normalized_messages = [
+                            collected_message
+                            for collected_message in normalized_messages
+                            if collected_message.type != ToolInvokeMessage.MessageType.UI
+                        ]
+                        logger.warning(
+                            "Ignored tool UI batch that exceeds agent invocation limits",
+                            exc_info=True,
+                        )
+                        continue
+
+                    ui_messages = candidate
+                    normalized_messages.append(normalized_message)
+            except (TypeError, ValueError):
+                if not is_ui_transport:
+                    raise
+                ui_batch_rejected = True
+                ui_messages.clear()
+                normalized_messages = [
+                    normalized_message
+                    for normalized_message in normalized_messages
+                    if normalized_message.type != ToolInvokeMessage.MessageType.UI
+                ]
+                logger.warning("Ignored invalid reserved tool UI message", exc_info=True)
+
+        return normalized_messages, ui_messages
+
+    @staticmethod
+    def _is_ui_transport_message(message: ToolInvokeMessage) -> bool:
+        if message.type == ToolInvokeMessage.MessageType.UI:
+            return True
+        if message.type == ToolInvokeMessage.MessageType.VARIABLE:
+            variable_message = cast(ToolInvokeMessage.VariableMessage, message.message)
+            return variable_message.variable_name == DIFY_UI_JSON_ENVELOPE_KEY
+        if message.type == ToolInvokeMessage.MessageType.JSON:
+            json_message = cast(ToolInvokeMessage.JsonMessage, message.message)
+            value = json_message.json_object
+            return isinstance(value, dict) and set(value) == {DIFY_UI_JSON_ENVELOPE_KEY}
+        return False
+
     @staticmethod
     def _extract_tool_response_binary_and_text(
         tool_response: list[ToolInvokeMessage],

api/core/workflow/node_runtime.py

@@ -90,6 +90,7 @@ from .system_variables import SystemVariableKey, get_system_text
 if TYPE_CHECKING:
     from core.tools.__base.tool import Tool
     from core.tools.entities.tool_entities import ToolInvokeMessage as CoreToolInvokeMessage
+    from core.tools.entities.ui_entities import ToolUIMessage
     from graphon.nodes.llm.file_saver import LLMFileSaver
     from graphon.nodes.tool.entities import ToolNodeData
 
@@ -666,6 +667,11 @@ class DifyToolNodeRuntime(ToolNodeRuntimeProtocol):
     ) -> Generator[ToolRuntimeMessage, None, None]:
         try:
             for message in messages:
+                # Tool Gen-UI is currently a chat presentation channel. Graphon
+                # does not define an equivalent runtime message, so workflow
+                # execution safely ignores it instead of coercing it to JSON.
+                if message.type.value == "ui":
+                    continue
                 yield self._convert_message(message)
         except Exception as exc:
             raise self._map_invocation_exception(exc, provider_name=provider_name) from exc
@@ -686,6 +692,7 @@ class DifyToolNodeRuntime(ToolNodeRuntimeProtocol):
         | CoreToolInvokeMessage.FileMessage
         | CoreToolInvokeMessage.VariableMessage
         | CoreToolInvokeMessage.RetrieverResourceMessage
+        | ToolUIMessage
         | None,
     ) -> (
         ToolRuntimeMessage.TextMessage

api/extensions/ext_login.py

@@ -1,5 +1,6 @@
 import json
-from typing import cast, override
+import logging
+from typing import assert_never, cast, override
 
 import flask_login
 from flask import Request, Response, request
@@ -11,6 +12,7 @@ from werkzeug.exceptions import NotFound, Unauthorized
 from configs import dify_config
 from constants import HEADER_NAME_APP_CODE
 from core.db.session_factory import session_factory
+from core.logging.context import set_identity_context
 from dify_app import DifyApp
 from libs.passport import PassportService
 from libs.token import extract_access_token, extract_console_cookie_token, extract_webapp_passport
@@ -19,6 +21,8 @@ from models.enums import EndUserType
 from models.model import AppMCPServer, EndUser
 from services.account_service import AccountService
 
+logger = logging.getLogger(__name__)
+
 type LoginUser = Account | EndUser
 
 
@@ -156,13 +160,24 @@ def _load_user_from_request(request_from_flask_login: Request, session: Session)
 @user_logged_in.connect
 @user_loaded_from_request.connect
 def on_user_logged_in(_sender: object, user: LoginUser) -> None:
-    """Called when a user logged in.
+    """Snapshot authenticated identity into the side-effect-free logging context.
 
     Note: AccountService.load_logged_in_account will populate user.current_tenant_id
     through the load_user method, which calls account.set_tenant_id_with_session().
     """
-    # tenant_id context variable removed - using current_user.current_tenant_id directly
-    pass
+    set_identity_context()
+    try:
+        match user:
+            case Account():
+                set_identity_context(tenant_id=user.current_tenant_id, user_id=user.id, user_type="account")
+            case EndUser():
+                set_identity_context(tenant_id=user.tenant_id, user_id=user.id, user_type=user.type or "end_user")
+            case _ as unreachable:
+                assert_never(unreachable)
+    except Exception:
+        # Logging enrichment must never make authentication fail.
+        logger.exception("Failed to set logging identity context")
+        return
 
 
 @login_manager.unauthorized_handler

api/extensions/otel/runtime.py

@@ -69,12 +69,17 @@ def on_user_loaded(_sender, user: Union["Account", "EndUser"]):
         if user:
             try:
                 current_span = get_current_span()
+                if not current_span.is_recording():
+                    return
                 tenant_id = extract_tenant_id(user)
                 if not tenant_id:
                     return
-                if current_span:
-                    current_span.set_attribute(DifySpanAttributes.TENANT_ID, tenant_id)
-                    current_span.set_attribute(GenAIAttributes.USER_ID, user.id)
+                current_span.set_attributes(
+                    {
+                        DifySpanAttributes.TENANT_ID: tenant_id,
+                        GenAIAttributes.USER_ID: user.id,
+                    }
+                )
             except Exception:
                 logger.exception("Error setting tenant and user attributes")
                 pass

api/fields/message_fields.py

@@ -7,6 +7,7 @@ from uuid import uuid4
 from pydantic import Field, computed_field, field_validator
 
 from core.entities.execution_extra_content import ExecutionExtraContentDomainModel
+from core.tools.entities.ui_entities import MessageUIPart, parse_history_ui_parts
 from fields.base import ResponseModel
 from fields.conversation_fields import AgentThought, JSONValue, MessageFile
 from graphon.file import File
@@ -64,6 +65,7 @@ class MessageListItem(ResponseModel):
     status: str
     error: str | None = None
     extra_contents: list[ExecutionExtraContentDomainModel]
+    ui_parts: list[MessageUIPart] = Field(default_factory=list, validation_alias="message_metadata_dict")
 
     @computed_field
     def total_tokens(self) -> int:
@@ -79,6 +81,12 @@ class MessageListItem(ResponseModel):
     def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
         return to_timestamp(value)
 
+    @field_validator("ui_parts", mode="before")
+    @classmethod
+    def _normalize_ui_parts(cls, value: object) -> list[MessageUIPart]:
+        raw_parts = value.get("ui_parts", []) if isinstance(value, dict) else value
+        return parse_history_ui_parts(raw_parts)
+
 
 class WebMessageListItem(MessageListItem):
     metadata: JSONValueType | None = Field(

api/openapi/markdown/console-openapi.md

@@ -17318,12 +17318,6 @@ Default model entity.
 | tool_name | string |  | Yes |
 | type | string |  | Yes |
 
-#### DeploymentEdition
-
-| Name | Type | Description | Required |
-| ---- | ---- | ----------- | -------- |
-| DeploymentEdition | string |  |  |
-
 #### DismissNotificationPayload
 
 | Name | Type | Description | Required |
@@ -19565,6 +19559,7 @@ Coarse node-level status used by Inspector to pick a banner.
 | ---- | ---- | ----------- | -------- |
 | avatar | string |  | No |
 | email | string |  | Yes |
+| id | string |  | Yes |
 | interface_language | string |  | Yes |
 | name | string |  | Yes |
 | timezone | string |  | Yes |
@@ -22032,7 +22027,6 @@ Model class for provider system configuration response.
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
 | branding | [BrandingModel](#brandingmodel) |  | Yes |
-| deployment_edition | [DeploymentEdition](#deploymentedition) |  | Yes |
 | enable_app_deploy | boolean |  | Yes |
 | enable_change_email | boolean, <br>**Default:** true |  | Yes |
 | enable_collaboration_mode | boolean, <br>**Default:** true |  | Yes |

api/openapi/markdown/web-openapi.md

@@ -1058,12 +1058,6 @@ Button styles for user actions.
 | auto_generate | boolean | Automatically generate the conversation name. When `true`, the `name` field is ignored. | No |
 | name | string | Conversation name. Required when `auto_generate` is `false`. | No |
 
-#### DeploymentEdition
-
-| Name | Type | Description | Required |
-| ---- | ---- | ----------- | -------- |
-| DeploymentEdition | string |  |  |
-
 #### EmailCodeLoginSendPayload
 
 | Name | Type | Description | Required |
@@ -1573,7 +1567,6 @@ Default configuration for form inputs.
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
 | branding | [BrandingModel](#brandingmodel) |  | Yes |
-| deployment_edition | [DeploymentEdition](#deploymentedition) |  | Yes |
 | enable_app_deploy | boolean |  | Yes |
 | enable_change_email | boolean, <br>**Default:** true |  | Yes |
 | enable_collaboration_mode | boolean, <br>**Default:** true |  | Yes |

api/services/agent/observability_service.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import json
+from collections.abc import Mapping
 from dataclasses import dataclass
 from datetime import datetime
 from decimal import Decimal
+from enum import Enum
 from typing import Any
 
 import sqlalchemy as sa
@@ -11,6 +14,7 @@ from sqlalchemy.orm import aliased
 
 from configs import dify_config
 from core.app.entities.app_invoke_entities import InvokeFrom
+from graphon.enums import WorkflowNodeExecutionStatus
 from libs.helper import convert_datetime_to_date, escape_like_pattern, to_timestamp
 from models.agent import WorkflowAgentNodeBinding
 from models.enums import CreatorUserRole, MessageStatus
@@ -194,11 +198,12 @@ class AgentObservabilityService:
         self, *, app: App, agent_id: str, conversation_id: str, params: AgentLogQueryParams
     ) -> dict[str, Any]:
         source_filters = self.resolve_source_filters(params.sources)
-        rows: list[Message] = []
+        rows: list[dict[str, Any]] = []
         for source_filter in source_filters:
             if source_filter.kind in {"all", "webapp"}:
                 rows.extend(
-                    self._list_webapp_messages(
+                    self.serialize_log_message(message)
+                    for message in self._list_webapp_messages(
                         app=app,
                         conversation_id=conversation_id,
                         params=params,
@@ -216,18 +221,18 @@ class AgentObservabilityService:
                     )
                 )
 
-        deduped = {message.id: message for message in rows}
-        sort_column = Message.created_at if params.sort_by == "created_at" else Message.updated_at
+        deduped = {row["id"]: row for row in rows}
+        sort_key = "created_at" if params.sort_by == "created_at" else "updated_at"
         sorted_rows = sorted(
             deduped.values(),
-            key=lambda message: (getattr(message, sort_column.key), message.id),
+            key=lambda row: (row[sort_key] or 0, row["id"]),
             reverse=params.sort_order != "asc",
         )
         total = len(sorted_rows)
         start = (params.page - 1) * params.limit
         end = start + params.limit
         return {
-            "data": [self.serialize_log_message(message) for message in sorted_rows[start:end]],
+            "data": sorted_rows[start:end],
             "page": params.page,
             "limit": params.limit,
             "total": total,
@@ -284,22 +289,20 @@ class AgentObservabilityService:
         workflow_app = aliased(App)
         stmt = (
             select(
-                Conversation,
+                WorkflowNodeExecutionModel.id.label("node_execution_id"),
+                WorkflowNodeExecutionModel.title.label("node_title"),
+                WorkflowNodeExecutionModel.status.label("node_status"),
+                WorkflowNodeExecutionModel.created_by_role.label("node_created_by_role"),
+                WorkflowNodeExecutionModel.created_by.label("node_created_by"),
+                WorkflowNodeExecutionModel.created_at.label("node_created_at"),
+                WorkflowNodeExecutionModel.finished_at.label("node_finished_at"),
                 workflow_app,
                 WorkflowAgentNodeBinding.workflow_id,
                 WorkflowAgentNodeBinding.workflow_version,
                 WorkflowAgentNodeBinding.node_id,
-                func.count(sa.distinct(Message.id)).label("message_count"),
-                func.max(Message.created_at).label("created_at"),
-                func.max(Message.updated_at).label("updated_at"),
-                func.sum(sa.case((Message.status == MessageStatus.PAUSED, 1), else_=0)).label("paused_count"),
-                func.sum(
-                    sa.case((or_(Message.error.is_not(None), Message.status == MessageStatus.ERROR), 1), else_=0)
-                ).label("failed_count"),
             )
-            .select_from(Message)
-            .join(Conversation, Conversation.id == Message.conversation_id)
-            .join(WorkflowRun, WorkflowRun.id == Message.workflow_run_id)
+            .select_from(WorkflowNodeExecutionModel)
+            .join(WorkflowRun, WorkflowRun.id == WorkflowNodeExecutionModel.workflow_run_id)
             .join(
                 WorkflowAgentNodeBinding,
                 and_(
@@ -310,40 +313,32 @@ class AgentObservabilityService:
                     WorkflowAgentNodeBinding.workflow_version == WorkflowRun.version,
                 ),
             )
-            .join(
-                WorkflowNodeExecutionModel,
-                and_(
-                    WorkflowNodeExecutionModel.workflow_run_id == WorkflowRun.id,
-                    WorkflowNodeExecutionModel.node_id == WorkflowAgentNodeBinding.node_id,
-                ),
-            )
             .join(workflow_app, workflow_app.id == WorkflowAgentNodeBinding.app_id)
-            .where(Message.workflow_run_id.is_not(None), Conversation.app_id == WorkflowAgentNodeBinding.app_id)
-            .group_by(
-                Conversation.id,
-                workflow_app.id,
-                WorkflowAgentNodeBinding.workflow_id,
-                WorkflowAgentNodeBinding.workflow_version,
-                WorkflowAgentNodeBinding.node_id,
+            .where(
+                WorkflowNodeExecutionModel.tenant_id == app.tenant_id,
+                WorkflowNodeExecutionModel.app_id == WorkflowAgentNodeBinding.app_id,
+                WorkflowNodeExecutionModel.workflow_id == WorkflowAgentNodeBinding.workflow_id,
+                WorkflowNodeExecutionModel.node_id == WorkflowAgentNodeBinding.node_id,
             )
         )
-        stmt = self._apply_observability_filters(stmt, params=params, source_filter=source_filter)
+        stmt = self._apply_workflow_node_filters(stmt, params=params, workflow_app=workflow_app)
         stmt = self._apply_workflow_source_filter(stmt, source_filter)
         rows = list(self._session.execute(stmt).all())
         return [
-            self._serialize_conversation_log(
-                conversation=row[0],
-                message_count=row.message_count,
-                paused_count=row.paused_count,
-                failed_count=row.failed_count,
+            self._serialize_workflow_execution_log(
+                node_execution_id=row.node_execution_id,
+                title=row.node_title,
+                status=row.node_status,
+                created_by_role=row.node_created_by_role,
+                created_by=row.node_created_by,
+                created_at=row.node_created_at,
+                finished_at=row.node_finished_at,
                 source=self._serialize_workflow_source(
-                    app=row[1],
+                    app=row[7],
                     workflow_id=row.workflow_id,
                     workflow_version=row.workflow_version,
                     node_id=row.node_id,
                 ),
-                created_at=row.created_at,
-                updated_at=row.updated_at,
             )
             for row in rows
         ]
@@ -363,10 +358,12 @@ class AgentObservabilityService:
         conversation_id: str,
         params: AgentLogQueryParams,
         source_filter: AgentSourceFilter,
-    ) -> list[Message]:
+    ) -> list[dict[str, Any]]:
+        workflow_app = aliased(App)
         stmt = (
-            select(Message)
-            .join(WorkflowRun, WorkflowRun.id == Message.workflow_run_id)
+            select(WorkflowNodeExecutionModel)
+            .select_from(WorkflowNodeExecutionModel)
+            .join(WorkflowRun, WorkflowRun.id == WorkflowNodeExecutionModel.workflow_run_id)
             .join(
                 WorkflowAgentNodeBinding,
                 and_(
@@ -377,18 +374,23 @@ class AgentObservabilityService:
                     WorkflowAgentNodeBinding.workflow_version == WorkflowRun.version,
                 ),
             )
-            .join(
-                WorkflowNodeExecutionModel,
-                and_(
-                    WorkflowNodeExecutionModel.workflow_run_id == WorkflowRun.id,
-                    WorkflowNodeExecutionModel.node_id == WorkflowAgentNodeBinding.node_id,
-                ),
+            .join(workflow_app, workflow_app.id == WorkflowAgentNodeBinding.app_id)
+            .where(
+                WorkflowNodeExecutionModel.id == conversation_id,
+                WorkflowNodeExecutionModel.tenant_id == app.tenant_id,
+                WorkflowNodeExecutionModel.app_id == WorkflowAgentNodeBinding.app_id,
+                WorkflowNodeExecutionModel.workflow_id == WorkflowAgentNodeBinding.workflow_id,
+                WorkflowNodeExecutionModel.node_id == WorkflowAgentNodeBinding.node_id,
             )
-            .where(Message.conversation_id == conversation_id)
         )
-        stmt = self._apply_message_filters(stmt, params=params, source_filter=source_filter)
+        stmt = self._apply_workflow_node_filters(stmt, params=params, workflow_app=workflow_app)
         stmt = self._apply_workflow_source_filter(stmt, source_filter)
-        return list(self._session.scalars(stmt.order_by(Message.created_at.desc(), Message.id.desc())).all())
+        executions = list(
+            self._session.scalars(
+                stmt.order_by(WorkflowNodeExecutionModel.created_at.desc(), WorkflowNodeExecutionModel.id.desc())
+            ).all()
+        )
+        return [self.serialize_workflow_node_message(execution) for execution in executions]
 
     def _list_workflow_sources(self, *, app: App, agent_id: str) -> list[dict[str, Any]]:
         workflow_app = aliased(App)
@@ -443,6 +445,62 @@ class AgentObservabilityService:
             stmt = cls._apply_status_filter(stmt, params.statuses)
         return stmt
 
+    @classmethod
+    def _apply_workflow_node_filters(cls, stmt, *, params: AgentLogQueryParams, workflow_app):
+        if params.start:
+            stmt = stmt.where(WorkflowNodeExecutionModel.created_at >= params.start)
+        if params.end:
+            stmt = stmt.where(WorkflowNodeExecutionModel.created_at < params.end)
+        if params.keyword:
+            escaped_keyword = escape_like_pattern(params.keyword)
+            pattern = f"%{escaped_keyword}%"
+            stmt = stmt.where(
+                or_(
+                    WorkflowNodeExecutionModel.inputs.ilike(pattern, escape="\\"),
+                    WorkflowNodeExecutionModel.outputs.ilike(pattern, escape="\\"),
+                    WorkflowNodeExecutionModel.error.ilike(pattern, escape="\\"),
+                    WorkflowNodeExecutionModel.title.ilike(pattern, escape="\\"),
+                    workflow_app.name.ilike(pattern, escape="\\"),
+                )
+            )
+        if params.statuses:
+            stmt = cls._apply_workflow_node_status_filter(stmt, params.statuses)
+        return stmt
+
+    @staticmethod
+    def _apply_workflow_node_status_filter(stmt, statuses: tuple[str, ...]):
+        conditions = []
+        for status in statuses:
+            normalized = status.strip().lower()
+            if normalized in {"success", "normal"}:
+                conditions.append(WorkflowNodeExecutionModel.status == WorkflowNodeExecutionStatus.SUCCEEDED)
+            elif normalized in {"failed", "error"}:
+                conditions.append(
+                    WorkflowNodeExecutionModel.status.in_(
+                        (
+                            WorkflowNodeExecutionStatus.FAILED,
+                            WorkflowNodeExecutionStatus.EXCEPTION,
+                            WorkflowNodeExecutionStatus.STOPPED,
+                        )
+                    )
+                )
+            elif normalized == "paused":
+                conditions.append(
+                    WorkflowNodeExecutionModel.status.in_(
+                        (
+                            WorkflowNodeExecutionStatus.PAUSED,
+                            WorkflowNodeExecutionStatus.PENDING,
+                            WorkflowNodeExecutionStatus.RUNNING,
+                            WorkflowNodeExecutionStatus.RETRY,
+                        )
+                    )
+                )
+            else:
+                raise ValueError(f"Unsupported status: {status}")
+        if not conditions:
+            return stmt
+        return stmt.where(or_(*conditions))
+
     @staticmethod
     def _apply_workflow_source_filter(stmt, source_filter: AgentSourceFilter):
         if source_filter.app_id:
@@ -505,6 +563,148 @@ class AgentObservabilityService:
             "updated_at": to_timestamp(updated_at or conversation.updated_at),
         }
 
+    @classmethod
+    def _serialize_workflow_execution_log(
+        cls,
+        *,
+        node_execution_id: str,
+        title: str,
+        status: object,
+        created_by_role: object,
+        created_by: str,
+        created_at: datetime,
+        finished_at: datetime | None,
+        source: dict[str, Any],
+    ) -> dict[str, Any]:
+        created_by_role_value = cls._enum_value(created_by_role)
+        return {
+            "id": node_execution_id,
+            "conversation_id": node_execution_id,
+            "title": title,
+            "end_user_id": created_by if created_by_role_value == CreatorUserRole.END_USER.value else None,
+            "message_count": 1,
+            "user_rate": None,
+            "operation_rate": None,
+            "unread": False,
+            "source": source,
+            "status": cls._workflow_node_status(status),
+            "created_at": to_timestamp(created_at),
+            "updated_at": to_timestamp(finished_at or created_at),
+        }
+
+    @classmethod
+    def serialize_workflow_node_message(cls, node_execution: WorkflowNodeExecutionModel) -> dict[str, Any]:
+        inputs = cls._json_mapping(node_execution.inputs)
+        outputs = cls._json_mapping(node_execution.outputs)
+        metadata = cls._json_mapping(node_execution.execution_metadata)
+        agent_log = cls._mapping_value(metadata, "agent_log")
+        agent_backend = cls._mapping_value(agent_log, "agent_backend")
+        usage = cls._mapping_value(agent_backend, "usage")
+
+        prompt_tokens = cls._int_value(usage.get("prompt_tokens"))
+        completion_tokens = cls._int_value(usage.get("completion_tokens"))
+        total_tokens = cls._int_value(usage.get("total_tokens") or metadata.get("total_tokens"))
+        if not total_tokens:
+            total_tokens = prompt_tokens + completion_tokens
+        created_by_role = cls._enum_value(node_execution.created_by_role)
+
+        return {
+            "id": node_execution.id,
+            "message_id": node_execution.id,
+            "conversation_id": node_execution.id,
+            "query": cls._workflow_node_query(inputs, fallback=node_execution.title),
+            "answer": cls._workflow_node_answer(outputs),
+            "status": cls._workflow_node_status(node_execution.status),
+            "error": node_execution.error,
+            "from_end_user_id": (
+                node_execution.created_by if created_by_role == CreatorUserRole.END_USER.value else None
+            ),
+            "from_account_id": (
+                node_execution.created_by if created_by_role == CreatorUserRole.ACCOUNT.value else None
+            ),
+            "message_tokens": prompt_tokens,
+            "answer_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "total_price": str(usage.get("total_price") or metadata.get("total_price") or Decimal(0)),
+            "currency": str(usage.get("currency") or metadata.get("currency") or ""),
+            "latency": float(usage.get("latency") or node_execution.elapsed_time or 0),
+            "created_at": to_timestamp(node_execution.created_at),
+            "updated_at": to_timestamp(node_execution.finished_at or node_execution.created_at),
+        }
+
+    @staticmethod
+    def _json_mapping(value: object) -> Mapping[str, Any]:
+        if isinstance(value, Mapping):
+            return value
+        if not isinstance(value, str) or not value:
+            return {}
+        try:
+            parsed = json.loads(value)
+        except (TypeError, ValueError):
+            return {}
+        return parsed if isinstance(parsed, Mapping) else {}
+
+    @staticmethod
+    def _mapping_value(value: Mapping[str, Any], key: str) -> Mapping[str, Any]:
+        nested = value.get(key)
+        return nested if isinstance(nested, Mapping) else {}
+
+    @staticmethod
+    def _enum_value(value: object) -> str:
+        return str(value.value) if isinstance(value, Enum) else str(value)
+
+    @staticmethod
+    def _int_value(value: object) -> int:
+        if not isinstance(value, (str, int, float, Decimal)):
+            return 0
+        try:
+            return int(value)
+        except (TypeError, ValueError):
+            return 0
+
+    @classmethod
+    def _workflow_node_query(cls, inputs: Mapping[str, Any], *, fallback: str) -> str:
+        request_data = cls._mapping_value(inputs, "agent_backend_request")
+        composition = cls._mapping_value(request_data, "composition")
+        layers = composition.get("layers")
+        prompts: list[str] = []
+        if isinstance(layers, list):
+            for layer_name in ("workflow_node_job_prompt", "workflow_user_prompt"):
+                for layer in layers:
+                    if not isinstance(layer, Mapping) or layer.get("name") != layer_name:
+                        continue
+                    config = cls._mapping_value(layer, "config")
+                    user_prompt = config.get("user")
+                    if isinstance(user_prompt, str) and user_prompt.strip():
+                        prompts.append(user_prompt.strip())
+        return "\n\n".join(prompts) or fallback
+
+    @staticmethod
+    def _workflow_node_answer(outputs: Mapping[str, Any]) -> str:
+        for key in ("output", "text", "answer"):
+            value = outputs.get(key)
+            if isinstance(value, str):
+                return value
+        return json.dumps(outputs, ensure_ascii=False) if outputs else ""
+
+    @classmethod
+    def _workflow_node_status(cls, status: object) -> str:
+        value = cls._enum_value(status)
+        if value in {
+            WorkflowNodeExecutionStatus.FAILED.value,
+            WorkflowNodeExecutionStatus.EXCEPTION.value,
+            WorkflowNodeExecutionStatus.STOPPED.value,
+        }:
+            return "failed"
+        if value in {
+            WorkflowNodeExecutionStatus.PAUSED.value,
+            WorkflowNodeExecutionStatus.PENDING.value,
+            WorkflowNodeExecutionStatus.RUNNING.value,
+            WorkflowNodeExecutionStatus.RETRY.value,
+        }:
+            return "paused"
+        return "success"
+
     @staticmethod
     def _conversation_status(*, paused_count: int, failed_count: int) -> str:
         if paused_count:

api/services/agent/roster_service.py

@@ -660,15 +660,18 @@ class AgentRosterService:
         agent_id: str,
         account_id: str,
         draft_type: AgentConfigDraftType = AgentConfigDraftType.DEBUG_BUILD,
-        commit: bool = True,
     ) -> str:
         """Start a new scoped console conversation for the current Agent App editor.
 
         If this account already has a mapping for the requested draft surface, the previous
-        conversation is abandoned first: any ACTIVE conversation-owned Agent
-        runtime sessions for that old conversation are sent through best-effort
-        backend cleanup and then retired locally even when enqueueing fails. The
-        other draft surface is left untouched.
+        conversation is abandoned after the replacement mapping is committed: any ACTIVE
+        conversation-owned Agent runtime sessions for that old conversation are sent through
+        best-effort backend cleanup and then retired locally even when enqueueing fails. This
+        order prevents a failed database commit from retiring the still-current runtime session.
+        The other draft surface is left untouched.
+
+        A user and draft surface own one current mapping. If new-conversation requests overlap,
+        the last committed rotation becomes current and earlier response IDs cannot be continued.
         """
 
         agent = self._session.scalar(
@@ -691,6 +694,7 @@ class AgentRosterService:
             app_id=backing_app_id,
             account_id=account_id,
         )
+        previous_conversation: tuple[str, str] | None = None
         mapping = self._session.scalar(
             select(AgentDebugConversation).where(
                 AgentDebugConversation.tenant_id == tenant_id,
@@ -714,19 +718,22 @@ class AgentRosterService:
             previous_app_id = mapping.app_id
             previous_conversation_id = mapping.conversation_id
             if previous_conversation_id:
-                self._cleanup_debug_conversation_runtime_sessions(
-                    tenant_id=tenant_id,
-                    agent_id=agent_id,
-                    account_id=account_id,
-                    draft_type=draft_type,
-                    app_id=previous_app_id or backing_app_id,
-                    conversation_id=previous_conversation_id,
-                )
+                previous_conversation = (previous_app_id or backing_app_id, previous_conversation_id)
             mapping.app_id = backing_app_id
             mapping.conversation_id = conversation_id
         self._session.flush()
-        if commit:
-            self._session.commit()
+        self._session.commit()
+
+        if previous_conversation:
+            previous_app_id, previous_conversation_id = previous_conversation
+            self._cleanup_debug_conversation_runtime_sessions(
+                tenant_id=tenant_id,
+                agent_id=agent_id,
+                account_id=account_id,
+                draft_type=draft_type,
+                app_id=previous_app_id,
+                conversation_id=previous_conversation_id,
+            )
         return conversation_id
 
     def _cleanup_debug_conversation_runtime_sessions(
@@ -739,8 +746,8 @@ class AgentRosterService:
         app_id: str,
         conversation_id: str,
     ) -> None:
-        session_store = AgentAppRuntimeSessionStore()
         try:
+            session_store = AgentAppRuntimeSessionStore()
             stored_sessions = session_store.list_active_sessions_for_conversation(
                 tenant_id=tenant_id,
                 app_id=app_id,

api/services/agent_tool_inner_service.py

@@ -92,6 +92,7 @@ class AgentToolInnerService:
                     conversation_id=request.caller.conversation_id,
                 )
             )
+            transformed_messages, _ = ToolEngine.collect_agent_messages(transformed_messages)
         except ToolProviderNotFoundError as exc:
             raise AgentToolInnerServiceError(
                 error_code="agent_tool_declaration_not_found",

api/services/feature_service.py

@@ -75,12 +75,6 @@ class LicenseStatus(StrEnum):
     LOST = "lost"
 
 
-class DeploymentEdition(StrEnum):
-    COMMUNITY = "COMMUNITY"
-    ENTERPRISE = "ENTERPRISE"
-    CLOUD = "CLOUD"
-
-
 class LicenseModel(FeatureResponseModel):
     status: LicenseStatus = LicenseStatus.NONE
     expired_at: str = ""
@@ -168,7 +162,6 @@ class PluginManagerModel(FeatureResponseModel):
 
 
 class SystemFeatureModel(FeatureResponseModel):
-    deployment_edition: DeploymentEdition
     enable_app_deploy: bool = False
     sso_enforced_for_signin: bool = False
     sso_enforced_for_signin_protocol: str = ""
@@ -259,7 +252,7 @@ class FeatureService:
 
     @classmethod
     def get_system_features(cls, is_authenticated: bool = False) -> SystemFeatureModel:
-        system_features = SystemFeatureModel(deployment_edition=cls._resolve_deployment_edition())
+        system_features = SystemFeatureModel()
         system_features.rbac_enabled = dify_config.RBAC_ENABLED
 
         cls._fulfill_system_params_from_env(system_features)
@@ -279,14 +272,6 @@ class FeatureService:
 
         return system_features
 
-    @classmethod
-    def _resolve_deployment_edition(cls) -> DeploymentEdition:
-        if dify_config.EDITION == "CLOUD":
-            return DeploymentEdition.CLOUD
-        if dify_config.ENTERPRISE_ENABLED:
-            return DeploymentEdition.ENTERPRISE
-        return DeploymentEdition.COMMUNITY
-
     @classmethod
     def get_app_dsl_version(cls) -> str:
         return CURRENT_APP_DSL_VERSION
@@ -300,7 +285,6 @@ class FeatureService:
         system_features.is_allow_register = dify_config.ALLOW_REGISTER
         system_features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE
         system_features.is_email_setup = dify_config.MAIL_TYPE is not None and dify_config.MAIL_TYPE != ""
-        system_features.enable_change_email = dify_config.ENABLE_CHANGE_EMAIL
         system_features.enable_trial_app = dify_config.ENABLE_TRIAL_APP
         system_features.enable_explore_banner = dify_config.ENABLE_EXPLORE_BANNER
         system_features.enable_learn_app = dify_config.ENABLE_LEARN_APP

api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth_server.py

@@ -318,6 +318,7 @@ def test_oauth_account_successful_retrieval(
 
     assert response.status_code == 200
     assert response.get_json() == {
+        "id": account.id,
         "name": "Test User",
         "email": account.email,
         "avatar": "avatar_url",

api/tests/test_containers_integration_tests/services/test_feature_service.py

@@ -414,7 +414,6 @@ class TestFeatureService:
             mock_config.ENABLE_EMAIL_PASSWORD_LOGIN = True
             mock_config.ENABLE_SOCIAL_OAUTH_LOGIN = False
             mock_config.ENABLE_COLLABORATION_MODE = False
-            mock_config.ENABLE_CHANGE_EMAIL = True
             mock_config.ALLOW_REGISTER = True
             mock_config.ALLOW_CREATE_WORKSPACE = True
             mock_config.MAIL_TYPE = "smtp"
@@ -617,7 +616,6 @@ class TestFeatureService:
             mock_config.ENABLE_EMAIL_CODE_LOGIN = False
             mock_config.ENABLE_EMAIL_PASSWORD_LOGIN = True
             mock_config.ENABLE_SOCIAL_OAUTH_LOGIN = True
-            mock_config.ENABLE_CHANGE_EMAIL = True
             mock_config.ALLOW_REGISTER = False
             mock_config.ALLOW_CREATE_WORKSPACE = False
             mock_config.MAIL_TYPE = None

api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py

@@ -69,6 +69,20 @@ def _version_response(version_id: str = "version-1") -> dict:
     }
 
 
+def test_query_values_accepts_repeated_and_indexed_arrays() -> None:
+    app = Flask(__name__)
+
+    with app.test_request_context("/?sources=webapp:app-1&sources%5B1%5D=workflow:app-2&sources%5B0%5D=workflow:app-1"):
+        assert roster_controller._query_values("sources", "source") == [
+            "webapp:app-1",
+            "workflow:app-1",
+            "workflow:app-2",
+        ]
+
+    with app.test_request_context("/?source%5B0%5D=workflow:app-3"):
+        assert roster_controller._query_values("sources", "source") == ["workflow:app-3"]
+
+
 def _workflow_composer_response(**overrides) -> dict:
     response = {
         "variant": "workflow",
@@ -1530,8 +1544,35 @@ def test_drain_streaming_generate_response_raises_when_stream_ends_early() -> No
         completion_controller._drain_streaming_generate_response(response)
 
 
-def test_agent_chat_helper_forces_agent_streaming_and_external_trace(
-    app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str
+@pytest.mark.parametrize(
+    ("payload_extra", "expected_draft_type", "expected_start_new"),
+    [
+        ({}, AgentConfigDraftType.DRAFT, True),
+        ({"conversation_id": None}, AgentConfigDraftType.DRAFT, True),
+        ({"conversation_id": ""}, AgentConfigDraftType.DRAFT, True),
+        (
+            {"conversation_id": "00000000-0000-0000-0000-000000000001"},
+            AgentConfigDraftType.DRAFT,
+            False,
+        ),
+        ({"draft_type": "debug_build"}, AgentConfigDraftType.DEBUG_BUILD, False),
+        (
+            {
+                "draft_type": "debug_build",
+                "conversation_id": "00000000-0000-0000-0000-000000000001",
+            },
+            AgentConfigDraftType.DEBUG_BUILD,
+            False,
+        ),
+    ],
+)
+def test_agent_chat_helper_resolves_scoped_conversation_and_forces_streaming(
+    app: Flask,
+    monkeypatch: pytest.MonkeyPatch,
+    account_id: str,
+    payload_extra: dict[str, str | None],
+    expected_draft_type: AgentConfigDraftType,
+    expected_start_new: bool,
 ) -> None:
     app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", mode="agent")
     current_user = SimpleNamespace(id=account_id)
@@ -1543,7 +1584,7 @@ def test_agent_chat_helper_forces_agent_streaming_and_external_trace(
 
     def resolve_debug_conversation(**kwargs: object) -> str:
         captured["resolve_debug_conversation"] = kwargs
-        return "debug-conversation-1"
+        return "00000000-0000-0000-0000-000000000001"
 
     monkeypatch.setattr(completion_controller.AppGenerateService, "generate", generate)
     monkeypatch.setattr(
@@ -1555,7 +1596,8 @@ def test_agent_chat_helper_forces_agent_streaming_and_external_trace(
         completion_controller.helper, "compact_generate_response", lambda response: {"response": response}
     )
     with app.test_request_context(
-        json={"inputs": {}, "query": "hello", "response_mode": "streaming"}, headers={"X-Trace-Id": "trace-1"}
+        json={"inputs": {}, "query": "hello", "response_mode": "streaming", **payload_extra},
+        headers={"X-Trace-Id": "trace-1"},
     ):
         result = completion_controller._create_chat_message(
             current_user=current_user, app_model=app_model, session=Mock()
@@ -1566,10 +1608,12 @@ def test_agent_chat_helper_forces_agent_streaming_and_external_trace(
     assert captured["streaming"] is True
     args = cast(dict[str, object], captured["args"])
     assert args["response_mode"] == "streaming"
-    assert args["conversation_id"] == "debug-conversation-1"
+    assert args["conversation_id"] == "00000000-0000-0000-0000-000000000001"
     assert args["auto_generate_name"] is False
     assert args["external_trace_id"] == "trace-1"
-    assert cast(dict[str, object], captured["resolve_debug_conversation"])["draft_type"] == AgentConfigDraftType.DRAFT
+    resolve_call = cast(dict[str, object], captured["resolve_debug_conversation"])
+    assert resolve_call["draft_type"] == expected_draft_type
+    assert resolve_call["start_new"] is expected_start_new
 
 
 def test_agent_chat_helper_ignores_private_exit_intent_payload_key(
@@ -1617,14 +1661,28 @@ def test_agent_chat_helper_ignores_private_exit_intent_payload_key(
     assert completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG not in args
 
 
-def test_agent_chat_helper_rejects_foreign_debug_conversation(
-    app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str
+@pytest.mark.parametrize(
+    ("payload_extra", "expected_draft_type"),
+    [
+        ({}, AgentConfigDraftType.DRAFT),
+        ({"draft_type": "debug_build"}, AgentConfigDraftType.DEBUG_BUILD),
+    ],
+)
+def test_agent_chat_helper_rejects_foreign_debug_conversation_before_generation(
+    app: Flask,
+    monkeypatch: pytest.MonkeyPatch,
+    account_id: str,
+    payload_extra: dict[str, str],
+    expected_draft_type: AgentConfigDraftType,
 ) -> None:
     app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", mode="agent")
+    generate = MagicMock()
+    resolve_debug_conversation = MagicMock(return_value="owned-conversation")
+    monkeypatch.setattr(completion_controller.AppGenerateService, "generate", generate)
     monkeypatch.setattr(
         completion_controller,
         "_resolve_current_user_agent_debug_conversation_id",
-        lambda **kwargs: "owned-conversation",
+        resolve_debug_conversation,
     )
     with app.test_request_context(
         json={
@@ -1632,6 +1690,7 @@ def test_agent_chat_helper_rejects_foreign_debug_conversation(
             "query": "hello",
             "response_mode": "streaming",
             "conversation_id": "00000000-0000-0000-0000-000000000001",
+            **payload_extra,
         }
     ):
         with pytest.raises(NotFound):
@@ -1643,6 +1702,12 @@ def test_agent_chat_helper_rejects_foreign_debug_conversation(
                 session=Mock(),
             )
 
+    resolve_debug_conversation.assert_called_once()
+    resolve_call = resolve_debug_conversation.call_args.kwargs
+    assert resolve_call["draft_type"] == expected_draft_type
+    assert resolve_call["start_new"] is False
+    generate.assert_not_called()
+
 
 def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app(
     monkeypatch: pytest.MonkeyPatch,
@@ -1657,6 +1722,10 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app
             calls.append({"get_or_create": kwargs})
             return f"debug-{kwargs['agent_id']}"
 
+        def refresh_agent_app_debug_conversation_id(self, **kwargs: object) -> str:
+            calls.append({"refresh": kwargs})
+            return f"new-{kwargs['agent_id']}"
+
         def get_app_backing_agent(self, **kwargs: object) -> object:
             calls.append({"get_app_backing_agent": kwargs})
             return SimpleNamespace(id="backing-agent")
@@ -1669,6 +1738,7 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app
         app_model=SimpleNamespace(id="app-1"),
         agent_id="agent-1",
         draft_type=AgentConfigDraftType.DRAFT,
+        start_new=True,
     )
     fallback_id = completion_controller._resolve_current_user_agent_debug_conversation_id(
         session="session-1",  # type: ignore[arg-type]
@@ -1678,10 +1748,20 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app
         agent_id=None,
         draft_type=AgentConfigDraftType.DEBUG_BUILD,
     )
-    assert explicit_id == "debug-agent-1"
+    fallback_preview_id = completion_controller._resolve_current_user_agent_debug_conversation_id(
+        session="session-1",  # type: ignore[arg-type]
+        current_tenant_id="tenant-1",
+        current_user=SimpleNamespace(id="account-1"),
+        app_model=SimpleNamespace(id="app-1"),
+        agent_id=None,
+        draft_type=AgentConfigDraftType.DRAFT,
+        start_new=True,
+    )
+    assert explicit_id == "new-agent-1"
     assert fallback_id == "debug-backing-agent"
+    assert fallback_preview_id == "new-backing-agent"
     assert calls[1] == {
-        "get_or_create": {
+        "refresh": {
             "tenant_id": "tenant-1",
             "agent_id": "agent-1",
             "account_id": "account-1",
@@ -1697,6 +1777,15 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app
             "draft_type": AgentConfigDraftType.DEBUG_BUILD,
         }
     }
+    assert calls[6] == {"get_app_backing_agent": {"tenant_id": "tenant-1", "app_id": "app-1"}}
+    assert calls[7] == {
+        "refresh": {
+            "tenant_id": "tenant-1",
+            "agent_id": "backing-agent",
+            "account_id": "account-1",
+            "draft_type": AgentConfigDraftType.DRAFT,
+        }
+    }
 
 
 @pytest.mark.parametrize(

api/tests/unit_tests/controllers/console/app/test_app_import_api.py

@@ -19,7 +19,7 @@ from models.engine import db
 from models.model import App, AppMode
 from services.app_dsl_service import ImportStatus
 from services.entities.dsl_entities import CheckDependenciesResult
-from services.feature_service import DeploymentEdition, SystemFeatureModel, WebAppAuthModel
+from services.feature_service import SystemFeatureModel, WebAppAuthModel
 
 
 def _unwrap(func):
@@ -47,10 +47,7 @@ class _Result:
 
 
 def _install_features(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
-    features = SystemFeatureModel(
-        deployment_edition=DeploymentEdition.COMMUNITY,
-        webapp_auth=WebAppAuthModel(enabled=enabled),
-    )
+    features = SystemFeatureModel(webapp_auth=WebAppAuthModel(enabled=enabled))
     monkeypatch.setattr(app_import_module.FeatureService, "get_system_features", lambda: features)
 
 

api/tests/unit_tests/controllers/console/auth/test_email_register.py

@@ -11,7 +11,7 @@ from controllers.console.auth.email_register import (
     EmailRegisterResetApi,
     EmailRegisterSendEmailApi,
 )
-from services.feature_service import DeploymentEdition, SystemFeatureModel
+from services.feature_service import SystemFeatureModel
 
 
 class TestEmailRegisterSendEmailApi:
@@ -34,11 +34,7 @@ class TestEmailRegisterSendEmailApi:
         mock_account = MagicMock()
         mock_get_account.return_value = mock_account
 
-        feature_flags = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-            is_allow_register=True,
-        )
+        feature_flags = SystemFeatureModel(enable_email_password_login=True, is_allow_register=True)
         with (
             patch("controllers.console.auth.email_register.dify_config.BILLING_ENABLED", True),
             patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
@@ -79,11 +75,7 @@ class TestEmailRegisterCheckApi:
         mock_get_data.return_value = {"email": "User@Example.com", "code": "4321"}
         mock_generate_token.return_value = (None, "new-token")
 
-        feature_flags = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-            is_allow_register=True,
-        )
+        feature_flags = SystemFeatureModel(enable_email_password_login=True, is_allow_register=True)
         with (
             patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
             patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags),
@@ -131,11 +123,7 @@ class TestEmailRegisterResetApi:
         mock_login.return_value = token_pair
         mock_get_account.return_value = None
 
-        feature_flags = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-            is_allow_register=True,
-        )
+        feature_flags = SystemFeatureModel(enable_email_password_login=True, is_allow_register=True)
         with (
             patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
             patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags),
@@ -183,11 +171,7 @@ class TestEmailRegisterResetApi:
         mock_login.return_value = token_pair
         mock_get_account.return_value = None
 
-        feature_flags = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-            is_allow_register=True,
-        )
+        feature_flags = SystemFeatureModel(enable_email_password_login=True, is_allow_register=True)
         with (
             patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
             patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags),
@@ -240,11 +224,7 @@ class TestEmailRegisterResetApi:
         mock_login.return_value = token_pair
         mock_get_account.return_value = None
 
-        feature_flags = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-            is_allow_register=True,
-        )
+        feature_flags = SystemFeatureModel(enable_email_password_login=True, is_allow_register=True)
         with (
             patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
             patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags),

api/tests/unit_tests/controllers/console/auth/test_forgot_password.py

@@ -15,7 +15,7 @@ from controllers.console.auth.forgot_password import (
 )
 from models.account import Account
 from models.engine import db
-from services.feature_service import DeploymentEdition, SystemFeatureModel
+from services.feature_service import SystemFeatureModel
 
 
 @pytest.fixture
@@ -46,15 +46,8 @@ class TestForgotPasswordSendEmailApi:
         mock_get_account.return_value = mock_account
         mock_send_email.return_value = "token-123"
 
-        wraps_features = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-            is_allow_register=True,
-        )
-        controller_features = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            is_allow_register=True,
-        )
+        wraps_features = SystemFeatureModel(enable_email_password_login=True, is_allow_register=True)
+        controller_features = SystemFeatureModel(is_allow_register=True)
         with (
             patch(
                 "controllers.console.auth.forgot_password.FeatureService.get_system_features",
@@ -102,10 +95,7 @@ class TestForgotPasswordCheckApi:
         mock_get_data.return_value = {"email": "Admin@Example.com", "code": "4321"}
         mock_generate_token.return_value = (None, "new-token")
 
-        wraps_features = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-        )
+        wraps_features = SystemFeatureModel(enable_email_password_login=True)
         with (
             patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
             patch("controllers.console.wraps.FeatureService.get_system_features", return_value=wraps_features),
@@ -148,10 +138,7 @@ class TestForgotPasswordResetApi:
         db.session.commit()
         mock_get_account.return_value = account
 
-        wraps_features = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-        )
+        wraps_features = SystemFeatureModel(enable_email_password_login=True)
         with (
             patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
             patch("controllers.console.wraps.FeatureService.get_system_features", return_value=wraps_features),

api/tests/unit_tests/controllers/console/auth/test_oauth_server.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from inspect import unwrap
 from unittest.mock import patch
 
-from controllers.console.auth.oauth_server import OAuthServerUserAuthorizeApi
+from controllers.console.auth.oauth_server import OAuthServerUserAccountApi, OAuthServerUserAuthorizeApi
 from models import Account
 from models.account import AccountStatus, TenantAccountRole
 from models.model import OAuthProviderApp
@@ -45,3 +45,13 @@ def test_oauth_authorize_uses_injected_current_user() -> None:
 
     sign_oauth_authorization_code.assert_called_once_with("client-1", "account-1")
     assert response == {"code": "authorization-code"}
+
+
+def test_oauth_account_returns_stable_account_id() -> None:
+    api = OAuthServerUserAccountApi()
+    method = unwrap(api.post)
+    account = _make_account()
+
+    response = method(api, _make_oauth_provider_app(), account)
+
+    assert response["id"] == "account-1"

api/tests/unit_tests/controllers/console/auth/test_password_reset.py

@@ -24,7 +24,7 @@ from controllers.console.auth.forgot_password import (
 )
 from controllers.console.error import AccountNotFound, EmailSendIpLimitError
 from models.account import Account, Tenant, TenantAccountJoin
-from services.feature_service import DeploymentEdition, SystemFeatureModel
+from services.feature_service import SystemFeatureModel
 
 SQLITE_MODELS = (Account, Tenant, TenantAccountJoin)
 
@@ -48,10 +48,7 @@ def enable_password_login_wrappers(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr("controllers.console.wraps.dify_config.EDITION", "CLOUD")
     monkeypatch.setattr(
         "controllers.console.wraps.FeatureService.get_system_features",
-        lambda: SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_email_password_login=True,
-        ),
+        lambda: SystemFeatureModel(enable_email_password_login=True),
     )
 
 

api/tests/unit_tests/controllers/console/test_feature.py

@@ -3,7 +3,7 @@ from inspect import unwrap
 from pytest_mock import MockerFixture
 
 from models import Account
-from services.feature_service import DeploymentEdition, FeatureModel, LimitationModel, SystemFeatureModel
+from services.feature_service import FeatureModel, LimitationModel, SystemFeatureModel
 
 
 def make_account() -> Account:
@@ -94,11 +94,7 @@ class TestSystemFeatureApi:
             "controllers.console.feature.current_account_with_tenant_optional",
             return_value=(account, "tenant-123"),
         )
-        system_features = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            is_allow_register=True,
-            enable_learn_app=True,
-        )
+        system_features = SystemFeatureModel(is_allow_register=True, enable_learn_app=True)
         get_system_features = mocker.patch(
             "controllers.console.feature.FeatureService.get_system_features",
             return_value=system_features,
@@ -123,10 +119,7 @@ class TestSystemFeatureApi:
             "controllers.console.feature.current_account_with_tenant_optional",
             return_value=(None, None),
         )
-        system_features = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            is_allow_register=False,
-        )
+        system_features = SystemFeatureModel(is_allow_register=False)
         get_system_features = mocker.patch(
             "controllers.console.feature.FeatureService.get_system_features",
             return_value=system_features,

api/tests/unit_tests/controllers/web/test_web_forgot_password.py

@@ -16,7 +16,7 @@ from controllers.web.forgot_password import (
 )
 from models.account import Account
 from models.engine import db
-from services.feature_service import DeploymentEdition, SystemFeatureModel
+from services.feature_service import SystemFeatureModel
 
 
 @pytest.fixture
@@ -32,10 +32,7 @@ def database_app() -> Iterator[Flask]:
 
 @pytest.fixture(autouse=True)
 def _patch_wraps():
-    wraps_features = SystemFeatureModel(
-        deployment_edition=DeploymentEdition.COMMUNITY,
-        enable_email_password_login=True,
-    )
+    wraps_features = SystemFeatureModel(enable_email_password_login=True)
     with (
         patch("controllers.console.wraps.db") as mock_db,
         patch("controllers.console.wraps.dify_config.ENTERPRISE_ENABLED", True),

api/tests/unit_tests/controllers/web/test_wraps.py

@@ -0,0 +1,49 @@
+from unittest import mock
+
+import pytest
+from werkzeug.exceptions import Unauthorized
+
+from core.logging.context import clear_request_context, get_identity_context
+
+
+@pytest.fixture(autouse=True)
+def _reset_logging_context():
+    clear_request_context()
+    yield
+    clear_request_context()
+
+
+def test_validate_jwt_token_sets_logging_identity_before_view() -> None:
+    from controllers.web import wraps
+
+    app_model = mock.Mock()
+    end_user = mock.Mock(id="end-user-id", tenant_id="tenant-id", type=None)
+    clear_request_context()
+
+    @wraps.validate_jwt_token
+    def protected_view(received_app, received_user):
+        assert get_identity_context() == ("tenant-id", "end-user-id", "end_user")
+        return received_app, received_user
+
+    with mock.patch.object(wraps, "decode_jwt_token", return_value=(app_model, end_user)):
+        result = protected_view()
+
+    assert result == (app_model, end_user)
+
+
+def test_validate_jwt_token_does_not_set_identity_when_authentication_fails() -> None:
+    from controllers.web import wraps
+
+    clear_request_context()
+
+    @wraps.validate_jwt_token
+    def protected_view(_app, _user):
+        raise AssertionError("view must not be called")
+
+    with (
+        mock.patch.object(wraps, "decode_jwt_token", side_effect=Unauthorized()),
+        pytest.raises(Unauthorized),
+    ):
+        protected_view()
+
+    assert get_identity_context() == ("", "", "")

api/tests/unit_tests/core/agent/test_base_agent_runner.py

@@ -8,6 +8,8 @@ from pytest_mock import MockerFixture
 import core.agent.base_agent_runner as module
 import models.model as model_module
 from core.agent.base_agent_runner import BaseAgentRunner
+from core.app.entities.queue_entities import QueueUIPartEvent
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID, ToolUIMessage, build_ui_part_id
 
 # ==========================================================
 # Fixtures
@@ -37,6 +39,45 @@ def runner(mocker: MockerFixture):
     return r
 
 
+def _ui_message(index: int) -> ToolUIMessage:
+    surface_id = f"surface-{index}"
+    return ToolUIMessage(
+        messages=[
+            {
+                "version": "v0.9.1",
+                "createSurface": {"surfaceId": surface_id, "catalogId": A2UI_CATALOG_ID},
+            },
+            {
+                "version": "v0.9.1",
+                "updateComponents": {
+                    "surfaceId": surface_id,
+                    "components": [{"id": "root", "component": "Text", "text": str(index)}],
+                },
+            },
+        ]
+    )
+
+
+def test_publish_ui_messages_enforces_batch_limit_before_queue(runner: BaseAgentRunner, mocker: MockerFixture) -> None:
+    runner.queue_manager = mocker.MagicMock()
+
+    runner.publish_ui_messages(
+        namespace="tool-call",
+        ui_messages=[_ui_message(index) for index in range(17)],
+    )
+
+    runner.queue_manager.publish.assert_not_called()
+
+    runner.publish_ui_messages(namespace="tool-call", ui_messages=[_ui_message(0)])
+    event = runner.queue_manager.publish.call_args.args[0]
+    assert isinstance(event, QueueUIPartEvent)
+    assert event.part.part_id == build_ui_part_id("tool-call", "surface-0")
+
+    runner.publish_ui_messages(namespace="x" * 10_000, ui_messages=[_ui_message(1)])
+    long_namespace_event = runner.queue_manager.publish.call_args.args[0]
+    assert len(long_namespace_event.part.part_id) <= 512
+
+
 # ==========================================================
 # _repack_app_generate_entity
 # ==========================================================

api/tests/unit_tests/core/agent/test_cot_agent_runner.py

@@ -10,6 +10,7 @@ from sqlalchemy.orm import Session
 from core.agent.cot_agent_runner import CotAgentRunner
 from core.agent.entities import AgentScratchpadUnit
 from core.agent.errors import AgentMaxIterationError
+from core.tools.tool_engine import ToolAgentInvokeResult
 from graphon.model_runtime.entities.llm_entities import LLMUsage
 
 
@@ -28,6 +29,19 @@ class DummyRunner(CotAgentRunner):
         return []
 
 
+def _invoke_result(
+    observation: str = "ok",
+    message_files: list[str] | None = None,
+    meta: MagicMock | None = None,
+) -> ToolAgentInvokeResult:
+    return ToolAgentInvokeResult(
+        observation=observation,
+        message_files=message_files or [],
+        ui_messages=[],
+        meta=meta or MagicMock(to_dict=lambda: {}),
+    )
+
+
 @pytest.fixture
 def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[DummyRunner]:
     # Prevent BaseAgentRunner __init__ from hitting database
@@ -181,7 +195,7 @@ class TestHandleInvokeAction:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=("result", [], MagicMock(to_dict=lambda: {})),
+            return_value=_invoke_result(observation="result"),
         )
 
         response, meta = runner._handle_invoke_action(runner.session, action, tool_instances, [])
@@ -226,7 +240,7 @@ class TestRun:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
+            return_value=_invoke_result(),
         )
 
         runner.agent_callback = None
@@ -248,7 +262,7 @@ class TestRun:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
+            return_value=_invoke_result(),
         )
 
         runner.agent_callback = None
@@ -323,7 +337,7 @@ class TestRun:
         runner.model_instance.invoke_llm = MagicMock(return_value=[])
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
+            return_value=_invoke_result(),
         )
 
         fake_prompt_tool = MagicMock()
@@ -391,7 +405,7 @@ class TestRun:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
+            return_value=_invoke_result(),
         )
 
         runner.app_config.agent.max_iteration = 5
@@ -452,7 +466,10 @@ class TestHandleInvokeActionExtended:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=("ok", ["file1"], MagicMock(to_dict=lambda: {"k": "v"})),
+            return_value=_invoke_result(
+                message_files=["file1"],
+                meta=MagicMock(to_dict=lambda: {"k": "v"}),
+            ),
         )
 
         message_file_ids = []

api/tests/unit_tests/core/agent/test_fc_agent_runner.py

@@ -12,6 +12,7 @@ from core.agent.errors import AgentMaxIterationError
 from core.agent.fc_agent_runner import FunctionCallAgentRunner
 from core.app.apps.base_app_queue_manager import PublishFrom
 from core.app.entities.queue_entities import QueueMessageFileEvent
+from core.tools.tool_engine import ToolAgentInvokeResult
 from graphon.model_runtime.entities.llm_entities import LLMUsage
 from graphon.model_runtime.entities.message_entities import (
     DocumentPromptMessageContent,
@@ -439,7 +440,12 @@ class TestRunMethod:
         tool_invoke_meta.to_dict.return_value = {"ok": True}
         mocker.patch(
             "core.agent.fc_agent_runner.ToolEngine.agent_invoke",
-            return_value=("ok", ["file1"], tool_invoke_meta),
+            return_value=ToolAgentInvokeResult(
+                observation="ok",
+                message_files=["file1"],
+                ui_messages=[],
+                meta=tool_invoke_meta,
+            ),
         )
 
         outputs = list(runner.run(runner.session, message, "query"))

api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py

@@ -55,7 +55,9 @@ from core.app.entities.queue_entities import (
     QueueAgentThoughtEvent,
     QueueLLMChunkEvent,
     QueueMessageEndEvent,
+    QueueUIPartEvent,
 )
+from core.tools.entities.ui_entities import build_ui_part_id
 from core.workflow.nodes.agent_v2.ask_human_resume import AskHumanResumeOutcome
 from graphon.model_runtime.errors.invoke import InvokeRateLimitError
 from models.agent_config_entities import AgentSoulConfig
@@ -1005,6 +1007,150 @@ def test_tool_call_part_binds_late_call_id_to_delta_row(monkeypatch):
     assert rows[0].observation == "Knowledge base search results: browser skill"
 
 
+def test_tool_return_metadata_publishes_validated_ui_part(monkeypatch):
+    fake_session = _FakeDbSession()
+    monkeypatch.setattr(app_runner_module.db, "session", fake_session)
+    qm = _FakeQueueManager()
+    recorder = app_runner_module._AgentProcessRecorder(
+        dify_context=_dify_ctx(),
+        message_id="msg-1",
+        queue_manager=qm,  # type: ignore[arg-type]
+    )
+    recorder.handle_stream_event(
+        AgentBackendStreamInternalEvent(
+            run_id="run-1",
+            data={
+                "event_kind": "function_tool_call",
+                "part": {
+                    "part_kind": "tool-call",
+                    "tool_name": "get_time",
+                    "args": {},
+                    "tool_call_id": "call-1",
+                },
+            },
+        )
+    )
+    recorder.handle_stream_event(
+        AgentBackendStreamInternalEvent(
+            run_id="run-1",
+            data={
+                "event_kind": "function_tool_result",
+                "part": {
+                    "part_kind": "tool-return",
+                    "tool_name": "get_time",
+                    "content": "It is 10:30.",
+                    "tool_call_id": "call-1",
+                    "metadata": {
+                        "dify_ui_messages": [
+                            {
+                                "protocol": "a2ui",
+                                "protocol_version": "v0.9.1",
+                                "messages": [
+                                    {
+                                        "version": "v0.9.1",
+                                        "createSurface": {
+                                            "surfaceId": "clock",
+                                            "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                                        },
+                                    },
+                                    {
+                                        "version": "v0.9.1",
+                                        "updateComponents": {
+                                            "surfaceId": "clock",
+                                            "components": [
+                                                {
+                                                    "id": "root",
+                                                    "component": "Text",
+                                                    "text": "10:30",
+                                                }
+                                            ],
+                                        },
+                                    },
+                                ],
+                            }
+                        ]
+                    },
+                },
+            },
+        )
+    )
+
+    ui_events = [event for event in qm.events if isinstance(event, QueueUIPartEvent)]
+    assert len(ui_events) == 1
+    assert ui_events[0].part.part_id == build_ui_part_id("call-1", "clock")
+    assert ui_events[0].part.sequence == 1
+    rows = list(fake_session.rows.values())
+    assert rows[0].observation == "It is 10:30."
+    assert "createSurface" not in rows[0].observation
+
+
+def test_tool_return_drops_oversized_ui_metadata_but_keeps_observation(monkeypatch):
+    fake_session = _FakeDbSession()
+    monkeypatch.setattr(app_runner_module.db, "session", fake_session)
+    qm = _FakeQueueManager()
+    recorder = app_runner_module._AgentProcessRecorder(
+        dify_context=_dify_ctx(),
+        message_id="msg-1",
+        queue_manager=qm,  # type: ignore[arg-type]
+    )
+    recorder.handle_stream_event(
+        AgentBackendStreamInternalEvent(
+            run_id="run-1",
+            data={
+                "event_kind": "function_tool_call",
+                "part": {
+                    "part_kind": "tool-call",
+                    "tool_name": "get_time",
+                    "args": {},
+                    "tool_call_id": "call-1",
+                },
+            },
+        )
+    )
+    ui_messages = [
+        {
+            "protocol": "a2ui",
+            "protocol_version": "v0.9.1",
+            "messages": [
+                {
+                    "version": "v0.9.1",
+                    "createSurface": {
+                        "surfaceId": f"clock-{index}",
+                        "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                    },
+                },
+                {
+                    "version": "v0.9.1",
+                    "updateComponents": {
+                        "surfaceId": f"clock-{index}",
+                        "components": [{"id": "root", "component": "Text", "text": str(index)}],
+                    },
+                },
+            ],
+        }
+        for index in range(17)
+    ]
+    recorder.handle_stream_event(
+        AgentBackendStreamInternalEvent(
+            run_id="run-1",
+            data={
+                "event_kind": "function_tool_result",
+                "part": {
+                    "part_kind": "tool-return",
+                    "tool_name": "get_time",
+                    "content": "It is 10:30.",
+                    "tool_call_id": "call-1",
+                    "metadata": {"dify_ui_messages": ui_messages},
+                },
+            },
+        )
+    )
+
+    assert not any(isinstance(event, QueueUIPartEvent) for event in qm.events)
+    rows = list(fake_session.rows.values())
+    assert rows[0].observation == "It is 10:30."
+
+
 def test_thinking_after_tool_starts_new_snapshot_row(monkeypatch):
     fake_session = _FakeDbSession()
     monkeypatch.setattr(app_runner_module.db, "session", fake_session)

api/tests/unit_tests/core/app/task_pipeline/test_ui_part_streaming.py

@@ -0,0 +1,131 @@
+from types import SimpleNamespace
+
+import pytest
+from pydantic import ValidationError
+
+from core.app.entities.queue_entities import QueueUIPartEvent
+from core.app.entities.task_entities import TaskStateMetadata
+from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID, MessageUIPart, ToolUIMessage
+
+
+def _part(sequence: int, text: str) -> MessageUIPart:
+    ui_message = ToolUIMessage.model_validate(
+        {
+            "protocol": "a2ui",
+            "protocol_version": "v0.9.1",
+            "messages": [
+                {
+                    "version": "v0.9.1",
+                    "createSurface": {
+                        "surfaceId": "clock",
+                        "catalogId": A2UI_CATALOG_ID,
+                    },
+                },
+                {
+                    "version": "v0.9.1",
+                    "updateComponents": {
+                        "surfaceId": "clock",
+                        "components": [{"id": "root", "component": "Text", "text": text}],
+                    },
+                },
+            ],
+        }
+    )
+    return MessageUIPart.from_tool_ui_message(
+        part_id="call-1:clock",
+        sequence=sequence,
+        ui_message=ui_message,
+    )
+
+
+def _distinct_part(index: int, *, large: bool = False) -> MessageUIPart:
+    surface_id = f"surface-{index}"
+    messages = [
+        {
+            "version": "v0.9.1",
+            "createSurface": {"surfaceId": surface_id, "catalogId": A2UI_CATALOG_ID},
+        }
+    ]
+    if large:
+        messages.append(
+            {
+                "version": "v0.9.1",
+                "updateDataModel": {
+                    "surfaceId": surface_id,
+                    "value": ["x" * 4096] * 20,
+                },
+            }
+        )
+    messages.append(
+        {
+            "version": "v0.9.1",
+            "updateComponents": {
+                "surfaceId": surface_id,
+                "components": [{"id": "root", "component": "Text", "text": str(index)}],
+            },
+        }
+    )
+    ui_message = ToolUIMessage(messages=messages)
+    return MessageUIPart.from_tool_ui_message(
+        part_id=f"call-{index}:{surface_id}",
+        sequence=1,
+        ui_message=ui_message,
+    )
+
+
+def test_ui_part_stream_response_upserts_only_newer_revision() -> None:
+    pipeline = object.__new__(EasyUIBasedGenerateTaskPipeline)
+    pipeline._application_generate_entity = SimpleNamespace(task_id="task-1")
+    pipeline._message_id = "message-1"
+    pipeline._task_state = SimpleNamespace(metadata=TaskStateMetadata())
+
+    first = pipeline._ui_part_to_stream_response(QueueUIPartEvent(part=_part(1, "10:30")))
+    stale = pipeline._ui_part_to_stream_response(QueueUIPartEvent(part=_part(1, "stale")))
+    newer = pipeline._ui_part_to_stream_response(QueueUIPartEvent(part=_part(2, "10:31")))
+
+    assert first is not None
+    assert first.id == "message-1"
+    assert first.part.sequence == 1
+    payload = first.model_dump(mode="json")
+    first_message = payload["part"]["messages"][0]
+    assert "createSurface" in first_message
+    assert "create_surface" not in first_message
+    assert first_message["createSurface"]["surfaceId"] == "clock"
+    assert first_message["createSurface"]["catalogId"] == A2UI_CATALOG_ID
+    component = payload["part"]["messages"][1]["updateComponents"]["components"][0]
+    assert component == {"id": "root", "component": "Text", "text": "10:30"}
+    assert stale is None
+    assert newer is not None
+    assert newer.part.sequence == 2
+    assert len(pipeline._task_state.metadata.ui_parts) == 1
+    assert pipeline._task_state.metadata.ui_parts[0].sequence == 2
+
+
+def test_ui_part_stream_response_drops_parts_over_message_budgets() -> None:
+    pipeline = object.__new__(EasyUIBasedGenerateTaskPipeline)
+    pipeline._application_generate_entity = SimpleNamespace(task_id="task-1")
+    pipeline._message_id = "message-1"
+    pipeline._task_state = SimpleNamespace(metadata=TaskStateMetadata())
+
+    responses = [
+        pipeline._ui_part_to_stream_response(QueueUIPartEvent(part=_distinct_part(index))) for index in range(17)
+    ]
+
+    assert all(response is not None for response in responses[:16])
+    assert responses[16] is None
+    assert len(pipeline._task_state.metadata.ui_parts) == 16
+    with pytest.raises(ValidationError, match="more than 16"):
+        TaskStateMetadata(ui_parts=[_distinct_part(index) for index in range(17)])
+
+    large_pipeline = object.__new__(EasyUIBasedGenerateTaskPipeline)
+    large_pipeline._application_generate_entity = SimpleNamespace(task_id="task-1")
+    large_pipeline._message_id = "message-2"
+    large_pipeline._task_state = SimpleNamespace(metadata=TaskStateMetadata())
+    large_responses = [
+        large_pipeline._ui_part_to_stream_response(QueueUIPartEvent(part=_distinct_part(index, large=True)))
+        for index in range(7)
+    ]
+
+    assert any(response is None for response in large_responses)
+    assert 0 < len(large_pipeline._task_state.metadata.ui_parts) < 7

api/tests/unit_tests/core/logging/test_context.py

@@ -1,15 +1,27 @@
 """Tests for logging context module."""
 
 import uuid
+from contextvars import copy_context
+
+import pytest
 
 from core.logging.context import (
     clear_request_context,
+    get_identity_context,
     get_request_id,
     get_trace_id,
     init_request_context,
+    set_identity_context,
 )
 
 
+@pytest.fixture(autouse=True)
+def _reset_logging_context():
+    clear_request_context()
+    yield
+    clear_request_context()
+
+
 class TestLoggingContext:
     """Tests for the logging context functions."""
 
@@ -77,3 +89,41 @@ class TestLoggingContext:
 
         # IDs should be different
         assert id1 != id2
+
+    def test_set_identity_context(self):
+        set_identity_context(tenant_id="tenant-1", user_id="user-1", user_type="end_user")
+
+        identity = get_identity_context()
+        assert identity.tenant_id == "tenant-1"
+        assert identity.user_id == "user-1"
+        assert identity.user_type == "end_user"
+
+    def test_set_identity_context_replaces_all_fields(self):
+        set_identity_context(tenant_id="tenant-1", user_id="user-1", user_type="account")
+
+        set_identity_context(user_id="user-2", user_type="end_user")
+
+        assert get_identity_context() == ("", "user-2", "end_user")
+
+    def test_identity_context_is_copied_as_primitive_values(self):
+        set_identity_context(tenant_id="tenant-1", user_id="user-1", user_type="end_user")
+        copied_context = copy_context()
+
+        clear_request_context()
+
+        assert get_identity_context() == ("", "", "")
+        assert copied_context.run(get_identity_context) == ("tenant-1", "user-1", "end_user")
+
+    def test_init_clears_existing_identity_context(self):
+        set_identity_context(tenant_id="tenant-1", user_id="user-1", user_type="end_user")
+
+        init_request_context()
+
+        assert get_identity_context() == ("", "", "")
+
+    def test_clear_resets_identity_context(self):
+        set_identity_context(tenant_id="tenant-1", user_id="user-1", user_type="end_user")
+
+        clear_request_context()
+
+        assert get_identity_context() == ("", "", "")

api/tests/unit_tests/core/logging/test_filters.py

@@ -1,5 +1,6 @@
 """Tests for logging filters."""
 
+import io
 import logging
 from unittest import mock
 
@@ -19,6 +20,15 @@ def log_record():
     )
 
 
+@pytest.fixture(autouse=True)
+def _reset_logging_context():
+    from core.logging.context import clear_request_context
+
+    clear_request_context()
+    yield
+    clear_request_context()
+
+
 class TestTraceContextFilter:
     def test_sets_empty_trace_id_without_context(self, log_record):
         from core.logging.context import clear_request_context
@@ -147,8 +157,10 @@ class TestTraceContextFilter:
 
 class TestIdentityContextFilter:
     def test_sets_empty_identity_without_request_context(self, log_record):
+        from core.logging.context import clear_request_context
         from core.logging.filters import IdentityContextFilter
 
+        clear_request_context()
         filter = IdentityContextFilter()
         result = filter.filter(log_record)
 
@@ -164,131 +176,92 @@ class TestIdentityContextFilter:
         result = filter.filter(log_record)
         assert result is True
 
-    def test_handles_exception_gracefully(self, log_record):
+    def test_uses_explicit_identity_context_without_flask_context(self, log_record):
+        from core.logging.context import set_identity_context
         from core.logging.filters import IdentityContextFilter
 
+        set_identity_context(tenant_id="tenant_id", user_id="end_user_id", user_type="end_user")
+
         filter = IdentityContextFilter()
+        filter.filter(log_record)
 
-        # Should not raise even if something goes wrong
-        with mock.patch(
-            "core.logging.filters.flask.has_request_context", side_effect=Exception("Test error"), autospec=True
-        ):
-            result = filter.filter(log_record)
-            assert result is True
-            assert log_record.tenant_id == ""
+        assert log_record.tenant_id == "tenant_id"
+        assert log_record.user_id == "end_user_id"
+        assert log_record.user_type == "end_user"
 
-    def test_sets_empty_identity_unauthenticated(self, log_record):
+    def test_does_not_trigger_flask_login_request_loader(self, log_record):
+        from flask import Flask
+        from flask_login import LoginManager
+
+        from core.logging.context import clear_request_context
         from core.logging.filters import IdentityContextFilter
 
-        mock_user = mock.MagicMock()
-        mock_user.is_authenticated = False
+        app = Flask(__name__)
+        app.secret_key = "test"
+        login_manager = LoginManager(app)
+        request_loader = mock.Mock(return_value=None)
+        login_manager.request_loader(request_loader)
+        clear_request_context()
 
-        with (
-            mock.patch("flask.has_request_context", return_value=True),
-            mock.patch("flask_login.current_user", mock_user),
-        ):
-            filter = IdentityContextFilter()
-            filter.filter(log_record)
-            assert log_record.user_id == ""
+        with app.test_request_context("/"):
+            from flask import g
 
-    def test_sets_identity_for_account(self, log_record):
+            assert "_login_user" not in g
+            IdentityContextFilter().filter(log_record)
+            assert "_login_user" not in g
+
+        request_loader.assert_not_called()
+        assert log_record.tenant_id == ""
+        assert log_record.user_id == ""
+        assert log_record.user_type == ""
+
+    def test_ended_otel_span_warning_does_not_trigger_request_loader(self):
+        from flask import Flask, g
+        from flask_login import LoginManager
+        from opentelemetry import trace
+        from opentelemetry.sdk.trace import TracerProvider
+
+        from core.logging.context import clear_request_context
         from core.logging.filters import IdentityContextFilter
 
-        class MockAccount:
-            pass
+        app = Flask(__name__)
+        app.secret_key = "test"
+        login_manager = LoginManager(app)
+        request_loader = mock.Mock(return_value=None)
+        login_manager.request_loader(request_loader)
 
-        mock_user = MockAccount()
-        mock_user.id = "account_id"
-        mock_user.current_tenant_id = "tenant_id"
-        mock_user.is_authenticated = True
+        span = TracerProvider().get_tracer(__name__).start_span("ended")
+        span.end()
 
-        with (
-            mock.patch("flask.has_request_context", return_value=True),
-            mock.patch("models.Account", MockAccount),
-            mock.patch("flask_login.current_user", mock_user),
-        ):
-            filter = IdentityContextFilter()
-            filter.filter(log_record)
+        stream = io.StringIO()
+        handler = logging.StreamHandler(stream)
+        handler.addFilter(IdentityContextFilter())
+        handler.setFormatter(logging.Formatter("%(tenant_id)s %(user_id)s %(user_type)s %(message)s"))
 
-            assert log_record.tenant_id == "tenant_id"
-            assert log_record.user_id == "account_id"
-            assert log_record.user_type == "account"
+        sdk_logger = logging.getLogger("opentelemetry.sdk.trace")
+        previous_level = sdk_logger.level
+        previous_propagate = sdk_logger.propagate
+        previous_disabled = sdk_logger.disabled
+        sdk_logger.addHandler(handler)
+        sdk_logger.setLevel(logging.WARNING)
+        sdk_logger.propagate = False
+        sdk_logger.disabled = False
+        clear_request_context()
 
-    def test_sets_identity_for_account_no_tenant(self, log_record):
-        from core.logging.filters import IdentityContextFilter
+        try:
+            with app.test_request_context("/"), trace.use_span(span, end_on_exit=False):
+                assert "_login_user" not in g
 
-        class MockAccount:
-            pass
+                span.set_attribute("test.key", "test-value")
 
-        mock_user = MockAccount()
-        mock_user.id = "account_id"
-        mock_user.current_tenant_id = None
-        mock_user.is_authenticated = True
+                assert "_login_user" not in g
+        finally:
+            clear_request_context()
+            sdk_logger.removeHandler(handler)
+            sdk_logger.setLevel(previous_level)
+            sdk_logger.propagate = previous_propagate
+            sdk_logger.disabled = previous_disabled
+            handler.close()
 
-        with (
-            mock.patch("flask.has_request_context", return_value=True),
-            mock.patch("models.Account", MockAccount),
-            mock.patch("flask_login.current_user", mock_user),
-        ):
-            filter = IdentityContextFilter()
-            filter.filter(log_record)
-
-            assert log_record.tenant_id == ""
-            assert log_record.user_id == "account_id"
-            assert log_record.user_type == "account"
-
-    def test_sets_identity_for_end_user(self, log_record):
-        from core.logging.filters import IdentityContextFilter
-
-        class MockEndUser:
-            pass
-
-        class AnotherClass:
-            pass
-
-        mock_user = MockEndUser()
-        mock_user.id = "end_user_id"
-        mock_user.tenant_id = "tenant_id"
-        mock_user.type = "custom_type"
-        mock_user.is_authenticated = True
-
-        with (
-            mock.patch("flask.has_request_context", return_value=True),
-            mock.patch("models.model.EndUser", MockEndUser),
-            mock.patch("models.Account", AnotherClass),
-            mock.patch("flask_login.current_user", mock_user),
-        ):
-            filter = IdentityContextFilter()
-            filter.filter(log_record)
-
-            assert log_record.tenant_id == "tenant_id"
-            assert log_record.user_id == "end_user_id"
-            assert log_record.user_type == "custom_type"
-
-    def test_sets_identity_for_end_user_default_type(self, log_record):
-        from core.logging.filters import IdentityContextFilter
-
-        class MockEndUser:
-            pass
-
-        class AnotherClass:
-            pass
-
-        mock_user = MockEndUser()
-        mock_user.id = "end_user_id"
-        mock_user.tenant_id = "tenant_id"
-        mock_user.type = None
-        mock_user.is_authenticated = True
-
-        with (
-            mock.patch("flask.has_request_context", return_value=True),
-            mock.patch("models.model.EndUser", MockEndUser),
-            mock.patch("models.Account", AnotherClass),
-            mock.patch("flask_login.current_user", mock_user),
-        ):
-            filter = IdentityContextFilter()
-            filter.filter(log_record)
-
-            assert log_record.tenant_id == "tenant_id"
-            assert log_record.user_id == "end_user_id"
-            assert log_record.user_type == "end_user"
+        request_loader.assert_not_called()
+        assert "Setting attribute on ended span" in stream.getvalue()

api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py

@@ -8,10 +8,49 @@ which provides document storage and retrieval functionality for datasets in the
 from unittest.mock import MagicMock, patch
 
 import pytest
+from sqlalchemy import func, select
+from sqlalchemy.orm import Session
 
 from core.rag.docstore.dataset_docstore import DatasetDocumentStore, DocumentSegment
-from core.rag.models.document import AttachmentDocument, Document
-from models.dataset import Dataset
+from core.rag.models.document import AttachmentDocument, ChildDocument, Document
+from models.dataset import ChildChunk, Dataset, SegmentAttachmentBinding
+
+TENANT_ID = "00000000-0000-0000-0000-000000000001"
+DATASET_ID = "00000000-0000-0000-0000-000000000002"
+DOCUMENT_ID = "00000000-0000-0000-0000-000000000003"
+USER_ID = "00000000-0000-0000-0000-000000000004"
+
+
+def _dataset() -> Dataset:
+    dataset = MagicMock(spec=Dataset)
+    dataset.id = DATASET_ID
+    dataset.tenant_id = TENANT_ID
+    return dataset
+
+
+def _persist_segment(
+    session: Session,
+    *,
+    index_node_id: str = "doc-1",
+    index_node_hash: str = "hash-1",
+    content: str = "Test content",
+    tokens: int = 5,
+) -> DocumentSegment:
+    segment = DocumentSegment(
+        tenant_id=TENANT_ID,
+        dataset_id=DATASET_ID,
+        document_id=DOCUMENT_ID,
+        position=1,
+        content=content,
+        word_count=len(content),
+        tokens=tokens,
+        created_by=USER_ID,
+        index_node_id=index_node_id,
+        index_node_hash=index_node_hash,
+    )
+    session.add(segment)
+    session.flush()
+    return segment
 
 
 class TestDatasetDocumentStoreInit:
@@ -132,228 +171,153 @@ class TestDatasetDocumentStoreDocs:
         assert result == {}
 
 
+@pytest.mark.parametrize(
+    "sqlite_session",
+    [(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
+    indirect=True,
+)
 class TestDatasetDocumentStoreAddDocuments:
     """Tests for add_documents method."""
 
-    def test_add_documents_new_document_with_embedding(self):
-        """Test adding new documents with embedding model."""
+    def test_add_documents_new_document_with_token_count(self, sqlite_session: Session):
+        """Test adding a new document with a precomputed token count."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_dataset.tenant_id = "tenant-1"
-        mock_dataset.indexing_technique = "high_quality"
-        mock_dataset.embedding_model_provider = "provider"
-        mock_dataset.embedding_model = "model"
+        document = Document(
+            page_content="Test content",
+            metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Test content"
-        mock_doc.metadata = {
-            "doc_id": "doc-1",
-            "doc_hash": "hash-1",
-        }
-        mock_doc.attachments = None
-        mock_doc.children = None
+        store.add_documents(session=sqlite_session, docs=[document], token_counts=[10])
+        sqlite_session.expire_all()
 
-        mock_model_instance = MagicMock()
-        mock_model_instance.get_text_embedding_num_tokens.return_value = [10]
+        segment = sqlite_session.scalar(
+            select(DocumentSegment).where(
+                DocumentSegment.dataset_id == DATASET_ID,
+                DocumentSegment.index_node_id == "doc-1",
+            )
+        )
+        assert segment is not None
+        assert segment.content == "Test content"
+        assert segment.tokens == 10
+        assert segment.position == 1
 
-        with (
-            patch("core.rag.docstore.dataset_docstore.ModelManager.for_tenant") as mock_manager_class,
-        ):
-            mock_session = MagicMock()
-            mock_session.scalar.return_value = None
-
-            mock_manager = MagicMock()
-            mock_manager.get_model_instance.return_value = mock_model_instance
-            mock_manager_class.return_value = mock_manager
-
-            with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
-                with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
-                    store = DatasetDocumentStore(
-                        dataset=mock_dataset,
-                        user_id="test-user-id",
-                        document_id="test-doc-id",
-                    )
-
-                    store.add_documents([mock_doc], session=mock_session)
-
-                    mock_session.add.assert_called()
-                    mock_session.flush.assert_called()
-
-    def test_add_documents_update_existing_document(self):
+    def test_add_documents_update_existing_document(self, sqlite_session: Session):
         """Test updating existing document with allow_update=True."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_dataset.tenant_id = "tenant-1"
-        mock_dataset.indexing_technique = "economy"
-        mock_dataset.embedding_model_provider = None
-        mock_dataset.embedding_model = None
+        existing_segment = _persist_segment(sqlite_session)
+        document = Document(
+            page_content="Updated content",
+            metadata={"doc_id": "doc-1", "doc_hash": "new-hash"},
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Updated content"
-        mock_doc.metadata = {
-            "doc_id": "doc-1",
-            "doc_hash": "new-hash",
-        }
-        mock_doc.attachments = None
-        mock_doc.children = None
+        store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
+        sqlite_session.expire_all()
 
-        mock_existing_segment = MagicMock()
-        mock_existing_segment.id = "seg-1"
+        updated_segment = sqlite_session.get(DocumentSegment, existing_segment.id)
+        assert updated_segment is not None
+        assert updated_segment.content == "Updated content"
+        assert updated_segment.index_node_hash == "new-hash"
+        assert updated_segment.tokens == 0
 
-        mock_session = MagicMock()
-        mock_session.scalar.return_value = 5
-
-        with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
-            with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
-                store = DatasetDocumentStore(
-                    dataset=mock_dataset,
-                    user_id="test-user-id",
-                    document_id="test-doc-id",
-                )
-
-                store.add_documents([mock_doc], session=mock_session)
-
-                mock_session.flush.assert_called()
-
-    def test_add_documents_raises_when_not_allowed(self):
+    def test_add_documents_raises_when_not_allowed(self, sqlite_session: Session):
         """Test that adding existing doc without allow_update raises ValueError."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_dataset.tenant_id = "tenant-1"
-        mock_dataset.indexing_technique = "economy"
+        _persist_segment(sqlite_session)
+        document = Document(
+            page_content="Test content",
+            metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Test content"
-        mock_doc.metadata = {
-            "doc_id": "doc-1",
-            "doc_hash": "hash-1",
-        }
-        mock_doc.attachments = None
-        mock_doc.children = None
-
-        mock_existing_segment = MagicMock()
-        mock_session = MagicMock()
-
-        with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
-            store = DatasetDocumentStore(
-                dataset=mock_dataset,
-                user_id="test-user-id",
-                document_id="test-doc-id",
+        with pytest.raises(ValueError, match="already exists"):
+            store.add_documents(
+                session=sqlite_session,
+                docs=[document],
+                token_counts=[0],
+                allow_update=False,
             )
 
-            with pytest.raises(ValueError, match="already exists"):
-                store.add_documents([mock_doc], session=mock_session, allow_update=False)
+        assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegment)) == 1
 
-    def test_add_documents_with_answer_metadata(self):
+    def test_add_documents_with_answer_metadata(self, sqlite_session: Session):
         """Test adding document with answer in metadata."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_dataset.tenant_id = "tenant-1"
-        mock_dataset.indexing_technique = "economy"
+        document = Document(
+            page_content="Test content",
+            metadata={
+                "doc_id": "doc-1",
+                "doc_hash": "hash-1",
+                "answer": "Test answer",
+            },
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Test content"
-        mock_doc.metadata = {
-            "doc_id": "doc-1",
-            "doc_hash": "hash-1",
-            "answer": "Test answer",
-        }
-        mock_doc.attachments = None
-        mock_doc.children = None
+        store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
+        sqlite_session.expire_all()
 
-        mock_session = MagicMock()
-        mock_session.scalar.return_value = None
+        segment = sqlite_session.scalar(select(DocumentSegment))
+        assert segment is not None
+        assert segment.answer == "Test answer"
 
-        with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
-            with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
-                store = DatasetDocumentStore(
-                    dataset=mock_dataset,
-                    user_id="test-user-id",
-                    document_id="test-doc-id",
-                )
-
-                store.add_documents([mock_doc], session=mock_session)
-
-                mock_session.add.assert_called()
-
-    def test_add_documents_with_invalid_document_type(self):
+    def test_add_documents_with_invalid_document_type(self, sqlite_session: Session):
         """Test that non-Document raises ValueError."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_session = MagicMock()
-
-        store = DatasetDocumentStore(
-            dataset=mock_dataset,
-            user_id="test-user-id",
-            document_id="test-doc-id",
-        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
         with pytest.raises(ValueError, match="must be a Document"):
-            store.add_documents(["not a document"], session=mock_session)
+            store.add_documents(session=sqlite_session, docs=["not a document"], token_counts=[0])  # type: ignore[list-item]
 
-    def test_add_documents_with_none_metadata(self):
+    def test_add_documents_with_none_metadata(self, sqlite_session: Session):
         """Test that document with None metadata raises ValueError."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Test content"
-        mock_doc.metadata = None
-        mock_session = MagicMock()
-
-        store = DatasetDocumentStore(
-            dataset=mock_dataset,
-            user_id="test-user-id",
-            document_id="test-doc-id",
-        )
+        document = MagicMock(spec=Document)
+        document.metadata = None
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
         with pytest.raises(ValueError, match="metadata must be a dict"):
-            store.add_documents([mock_doc], session=mock_session)
+            store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
 
-    def test_add_documents_with_save_child(self):
+    def test_add_documents_with_save_child(self, sqlite_session: Session):
         """Test adding documents with save_child=True."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_dataset.tenant_id = "tenant-1"
-        mock_dataset.indexing_technique = "economy"
-
-        mock_child = MagicMock(spec=Document)
-        mock_child.page_content = "Child content"
-        mock_child.metadata = {
-            "doc_id": "child-1",
-            "doc_hash": "child-hash",
-        }
-
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Test content"
-        mock_doc.metadata = {
-            "doc_id": "doc-1",
-            "doc_hash": "hash-1",
-        }
-        mock_doc.attachments = None
-        mock_doc.children = [mock_child]
-
-        mock_session = MagicMock()
-        mock_session.scalar.return_value = None
-
-        with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None):
-            with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
-                store = DatasetDocumentStore(
-                    dataset=mock_dataset,
-                    user_id="test-user-id",
-                    document_id="test-doc-id",
+        document = Document(
+            page_content="Test content",
+            metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
+            children=[
+                ChildDocument(
+                    page_content="Child content",
+                    metadata={"doc_id": "child-1", "doc_hash": "child-hash"},
                 )
+            ],
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
-                store.add_documents([mock_doc], session=mock_session, save_child=True)
+        store.add_documents(
+            session=sqlite_session,
+            docs=[document],
+            token_counts=[0],
+            save_child=True,
+        )
+        sqlite_session.expire_all()
 
-                mock_session.add.assert_called()
+        child = sqlite_session.scalar(select(ChildChunk))
+        assert child is not None
+        assert child.content == "Child content"
+        assert child.index_node_id == "child-1"
+
+    def test_add_documents_rejects_mismatched_token_counts(self, sqlite_session: Session):
+        document = Document(
+            page_content="Test content",
+            metadata={"doc_id": "doc-1", "doc_hash": "hash-1"},
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
+
+        with pytest.raises(ValueError):
+            store.add_documents(session=sqlite_session, docs=[document], token_counts=[])
+
+        assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegment)) == 0
 
 
 class TestDatasetDocumentStoreExists:
@@ -722,88 +686,85 @@ class TestDatasetDocumentStoreMultimodelBinding:
         mock_session.add.assert_not_called()
 
 
+@pytest.mark.parametrize(
+    "sqlite_session",
+    [(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
+    indirect=True,
+)
 class TestDatasetDocumentStoreAddDocumentsUpdateChild:
     """Tests for add_documents when updating existing documents with children."""
 
-    def test_add_documents_update_existing_with_children(self):
+    def test_add_documents_update_existing_with_children(self, sqlite_session: Session):
         """Test updating existing document with save_child=True and children."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_dataset.tenant_id = "tenant-1"
-        mock_dataset.indexing_technique = "economy"
-
-        mock_child = MagicMock(spec=Document)
-        mock_child.page_content = "Updated child content"
-        mock_child.metadata = {
-            "doc_id": "child-1",
-            "doc_hash": "new-child-hash",
-        }
-
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Updated content"
-        mock_doc.metadata = {
-            "doc_id": "doc-1",
-            "doc_hash": "new-hash",
-        }
-        mock_doc.attachments = None
-        mock_doc.children = [mock_child]
-
-        mock_existing_segment = MagicMock()
-        mock_existing_segment.id = "seg-1"
-
-        mock_session = MagicMock()
-        mock_session.scalar.return_value = 5
-
-        with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
-            with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
-                store = DatasetDocumentStore(
-                    dataset=mock_dataset,
-                    user_id="test-user-id",
-                    document_id="test-doc-id",
+        segment = _persist_segment(sqlite_session)
+        sqlite_session.add(
+            ChildChunk(
+                tenant_id=TENANT_ID,
+                dataset_id=DATASET_ID,
+                document_id=DOCUMENT_ID,
+                segment_id=segment.id,
+                position=1,
+                index_node_id="old-child",
+                index_node_hash="old-child-hash",
+                content="Old child content",
+                word_count=len("Old child content"),
+                created_by=USER_ID,
+            )
+        )
+        sqlite_session.flush()
+        document = Document(
+            page_content="Updated content",
+            metadata={"doc_id": "doc-1", "doc_hash": "new-hash"},
+            children=[
+                ChildDocument(
+                    page_content="Updated child content",
+                    metadata={"doc_id": "child-1", "doc_hash": "new-child-hash"},
                 )
+            ],
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
-                store.add_documents([mock_doc], session=mock_session, save_child=True)
+        store.add_documents(
+            session=sqlite_session,
+            docs=[document],
+            token_counts=[0],
+            save_child=True,
+        )
+        sqlite_session.expire_all()
 
-                mock_session.execute.assert_called()
-                mock_session.flush.assert_called()
+        children = sqlite_session.scalars(select(ChildChunk).order_by(ChildChunk.position)).all()
+        assert len(children) == 1
+        assert children[0].content == "Updated child content"
+        assert children[0].index_node_id == "child-1"
 
 
+@pytest.mark.parametrize(
+    "sqlite_session",
+    [(DocumentSegment, ChildChunk, SegmentAttachmentBinding)],
+    indirect=True,
+)
 class TestDatasetDocumentStoreAddDocumentsUpdateAnswer:
     """Tests for add_documents when updating existing documents with answer metadata."""
 
-    def test_add_documents_update_existing_with_answer(self):
+    def test_add_documents_update_existing_with_answer(self, sqlite_session: Session):
         """Test updating existing document with answer in metadata."""
 
-        mock_dataset = MagicMock(spec=Dataset)
-        mock_dataset.id = "test-dataset-id"
-        mock_dataset.tenant_id = "tenant-1"
-        mock_dataset.indexing_technique = "economy"
+        existing_segment = _persist_segment(sqlite_session)
+        document = Document(
+            page_content="Updated content",
+            metadata={
+                "doc_id": "doc-1",
+                "doc_hash": "new-hash",
+                "answer": "Updated answer",
+            },
+        )
+        store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID, document_id=DOCUMENT_ID)
 
-        mock_doc = MagicMock(spec=Document)
-        mock_doc.page_content = "Updated content"
-        mock_doc.metadata = {
-            "doc_id": "doc-1",
-            "doc_hash": "new-hash",
-            "answer": "Updated answer",
-        }
-        mock_doc.attachments = None
-        mock_doc.children = None
+        store.add_documents(session=sqlite_session, docs=[document], token_counts=[0])
+        sqlite_session.expire_all()
 
-        mock_existing_segment = MagicMock()
-        mock_existing_segment.id = "seg-1"
-
-        mock_session = MagicMock()
-        mock_session.scalar.return_value = 5
-
-        with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_existing_segment):
-            with patch.object(DatasetDocumentStore, "add_multimodel_documents_binding"):
-                store = DatasetDocumentStore(
-                    dataset=mock_dataset,
-                    user_id="test-user-id",
-                    document_id="test-doc-id",
-                )
-
-                store.add_documents([mock_doc], session=mock_session)
-
-                mock_session.flush.assert_called()
+        updated_segment = sqlite_session.get(DocumentSegment, existing_segment.id)
+        assert updated_segment is not None
+        assert updated_segment.answer == "Updated answer"
+        assert updated_segment.tokens == 0

api/tests/unit_tests/core/rag/embedding/test_token_counter.py

@@ -0,0 +1,56 @@
+from unittest.mock import Mock, patch
+
+from core.rag.embedding.token_counter import calculate_segment_token_counts
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
+from core.rag.models.document import Document
+from models.dataset import Dataset
+
+
+def test_high_quality_counts_each_document_once() -> None:
+    dataset = Mock(spec=Dataset)
+    dataset.tenant_id = "tenant-1"
+    dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
+    dataset.embedding_model_provider = "provider"
+    dataset.embedding_model = "model"
+    documents = [
+        Document(page_content="first", metadata={}),
+        Document(page_content="second", metadata={}),
+        Document(page_content="third", metadata={}),
+    ]
+
+    with patch("core.rag.embedding.token_counter.ModelManager.for_tenant") as model_manager_factory:
+        embedding_model = model_manager_factory.return_value.get_model_instance.return_value
+        embedding_model.get_text_embedding_num_tokens.return_value = [11, 22, 33]
+
+        result = calculate_segment_token_counts(dataset=dataset, documents=documents)
+
+    assert result == [11, 22, 33]
+    model_manager_factory.assert_called_once_with(tenant_id=dataset.tenant_id)
+    model_manager_factory.return_value.get_model_instance.assert_called_once()
+    embedding_model.get_text_embedding_num_tokens.assert_called_once_with(["first", "second", "third"])
+
+
+def test_economy_returns_zero_without_loading_model() -> None:
+    dataset = Mock(spec=Dataset)
+    dataset.indexing_technique = IndexTechniqueType.ECONOMY
+    documents = [
+        Document(page_content="first", metadata={}),
+        Document(page_content="second", metadata={}),
+    ]
+
+    with patch("core.rag.embedding.token_counter.ModelManager.for_tenant") as model_manager_factory:
+        result = calculate_segment_token_counts(dataset=dataset, documents=documents)
+
+    assert result == [0, 0]
+    model_manager_factory.assert_not_called()
+
+
+def test_empty_documents_return_without_loading_model() -> None:
+    dataset = Mock(spec=Dataset)
+    dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
+
+    with patch("core.rag.embedding.token_counter.ModelManager.for_tenant") as model_manager_factory:
+        result = calculate_segment_token_counts(dataset=dataset, documents=[])
+
+    assert result == []
+    model_manager_factory.assert_not_called()

api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py

@@ -271,14 +271,26 @@ class TestParagraphIndexProcessor:
             patch(
                 "core.rag.index_processor.processor.paragraph_index_processor.DatasetDocumentStore"
             ) as mock_store_cls,
+            patch(
+                "core.rag.index_processor.processor.paragraph_index_processor.calculate_segment_token_counts"
+            ) as mock_token_counter,
             patch("core.rag.index_processor.processor.paragraph_index_processor.Vector") as mock_vector_cls,
         ):
+            mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [11, 22]
             mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
             mock_vector_cls.return_value.create.side_effect = lambda _documents: phase_events.append("vector")
             processor.index(dataset, dataset_document, ["chunk-1", "chunk-2"], session)
 
-        assert phase_events == ["store", "commit", "vector"]
-        mock_store_cls.return_value.add_documents.assert_called_once()
+        assert phase_events == ["count", "store", "commit", "vector"]
+        documents = mock_token_counter.call_args.kwargs["documents"]
+        assert [document.page_content for document in documents] == ["chunk-1", "chunk-2"]
+        mock_token_counter.assert_called_once_with(dataset=dataset, documents=documents)
+        mock_store_cls.return_value.add_documents.assert_called_once_with(
+            session=session,
+            docs=documents,
+            token_counts=[11, 22],
+            save_child=False,
+        )
         mock_vector_cls.assert_called_once_with(dataset, session=session)
         mock_vector_cls.return_value.create.assert_called_once()
         mock_vector_cls.return_value.create_multimodal.assert_called_once()
@@ -299,13 +311,18 @@ class TestParagraphIndexProcessor:
             patch(
                 "core.rag.index_processor.processor.paragraph_index_processor.DatasetDocumentStore"
             ) as mock_store_cls,
+            patch(
+                "core.rag.index_processor.processor.paragraph_index_processor.calculate_segment_token_counts"
+            ) as mock_token_counter,
             patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls,
         ):
+            mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [0]
             mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
             mock_keyword_cls.return_value.add_texts.side_effect = lambda *_args: phase_events.append("keyword")
             processor.index(dataset, dataset_document, ["chunk-3"], session)
 
-        assert phase_events == ["store", "commit", "keyword"]
+        assert phase_events == ["count", "store", "commit", "keyword"]
+        mock_token_counter.assert_called_once()
         mock_keyword_cls.return_value.add_texts.assert_called_once()
 
     def test_index_multimodal_structure_handles_files_and_account_lookup(
@@ -341,6 +358,10 @@ class TestParagraphIndexProcessor:
                 processor, "_get_content_files", return_value=[AttachmentDocument(page_content="img", metadata={})]
             ) as mock_files,
             patch("core.rag.index_processor.processor.paragraph_index_processor.DatasetDocumentStore"),
+            patch(
+                "core.rag.index_processor.processor.paragraph_index_processor.calculate_segment_token_counts",
+                return_value=[11, 22],
+            ),
             patch("core.rag.index_processor.processor.paragraph_index_processor.Vector"),
         ):
             processor.index(dataset, dataset_document, {"general_chunks": []}, session)

api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py

@@ -362,17 +362,29 @@ class TestParentChildIndexProcessor:
             patch(
                 "core.rag.index_processor.processor.parent_child_index_processor.DatasetDocumentStore"
             ) as mock_store_cls,
+            patch(
+                "core.rag.index_processor.processor.parent_child_index_processor.calculate_segment_token_counts"
+            ) as mock_token_counter,
             patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls,
         ):
+            mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [11]
             mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
             mock_vector_cls.return_value.create.side_effect = lambda _documents: phase_events.append("vector")
             processor.index(dataset, dataset_document, {"parent_child_chunks": []}, session)
 
-        assert phase_events == ["store", "commit", "vector"]
+        assert phase_events == ["count", "store", "commit", "vector"]
         assert dataset_document.dataset_process_rule_id == "rule-1"
         session.add.assert_called_once_with(dataset_rule)
         session.flush.assert_called_once()
-        mock_store_cls.return_value.add_documents.assert_called_once()
+        documents = mock_token_counter.call_args.kwargs["documents"]
+        assert [document.page_content for document in documents] == ["parent text"]
+        mock_token_counter.assert_called_once_with(dataset=dataset, documents=documents)
+        mock_store_cls.return_value.add_documents.assert_called_once_with(
+            session=session,
+            docs=documents,
+            token_counts=[11],
+            save_child=True,
+        )
         mock_vector_cls.assert_called_once_with(dataset, session=session)
         assert mock_vector_cls.return_value.create.call_count == 1
         mock_vector_cls.return_value.create_multimodal.assert_called_once()
@@ -413,6 +425,10 @@ class TestParentChildIndexProcessor:
                 processor, "_get_content_files", return_value=[AttachmentDocument(page_content="image", metadata={})]
             ) as mock_files,
             patch("core.rag.index_processor.processor.parent_child_index_processor.DatasetDocumentStore"),
+            patch(
+                "core.rag.index_processor.processor.parent_child_index_processor.calculate_segment_token_counts",
+                return_value=[11],
+            ),
             patch("core.rag.index_processor.processor.parent_child_index_processor.Vector"),
         ):
             processor.index(dataset, dataset_document, {"parent_child_chunks": []}, session)

api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py

@@ -292,14 +292,26 @@ class TestQAIndexProcessor:
                 "core.rag.index_processor.processor.qa_index_processor.helper.generate_text_hash", return_value="hash"
             ),
             patch("core.rag.index_processor.processor.qa_index_processor.DatasetDocumentStore") as mock_store_cls,
+            patch(
+                "core.rag.index_processor.processor.qa_index_processor.calculate_segment_token_counts"
+            ) as mock_token_counter,
             patch("core.rag.index_processor.processor.qa_index_processor.Vector") as mock_vector_cls,
         ):
+            mock_token_counter.side_effect = lambda **_kwargs: phase_events.append("count") or [11, 22]
             mock_store_cls.return_value.add_documents.side_effect = lambda **_kwargs: phase_events.append("store")
             mock_vector_cls.return_value.create.side_effect = lambda _documents: phase_events.append("vector")
             processor.index(dataset, dataset_document, {"qa_chunks": []}, session)
 
-        assert phase_events == ["store", "commit", "vector"]
-        mock_store_cls.return_value.add_documents.assert_called_once()
+        assert phase_events == ["count", "store", "commit", "vector"]
+        documents = mock_token_counter.call_args.kwargs["documents"]
+        assert [document.page_content for document in documents] == ["Q1", "Q2"]
+        mock_token_counter.assert_called_once_with(dataset=dataset, documents=documents)
+        mock_store_cls.return_value.add_documents.assert_called_once_with(
+            session=session,
+            docs=documents,
+            token_counts=[11, 22],
+            save_child=False,
+        )
         mock_vector_cls.return_value.create.assert_called_once()
 
     def test_index_requires_high_quality(
@@ -318,6 +330,10 @@ class TestQAIndexProcessor:
                 "core.rag.index_processor.processor.qa_index_processor.helper.generate_text_hash", return_value="hash"
             ),
             patch("core.rag.index_processor.processor.qa_index_processor.DatasetDocumentStore"),
+            patch(
+                "core.rag.index_processor.processor.qa_index_processor.calculate_segment_token_counts",
+                return_value=[0],
+            ),
         ):
             with pytest.raises(ValueError, match="must be high quality"):
                 processor.index(dataset, dataset_document, {"qa_chunks": []}, session)

api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py

@@ -69,6 +69,7 @@ from graphon.model_runtime.entities.model_entities import ModelType
 from libs.datetime_utils import naive_utc_now
 from models.dataset import Dataset, DatasetProcessRule, DocumentSegment
 from models.dataset import Document as DatasetDocument
+from models.enums import SegmentStatus
 from models.model import Account
 
 # ============================================================================
@@ -611,7 +612,7 @@ class TestIndexingRunnerLoad:
     - Keyword index creation
     - Multi-threaded processing
     - Document segment status updates
-    - Token counting
+    - Precomputed token totals
     - Error handling during loading
     """
 
@@ -677,16 +678,10 @@ class TestIndexingRunnerLoad:
         """Test loading with high quality indexing (vector embeddings)."""
         # Arrange
         runner = IndexingRunner()
-        mock_embedding_instance = MagicMock()
-        mock_embedding_instance.get_text_embedding_num_tokens.return_value = 100
-        model_manager = mock_dependencies["model_manager"].return_value
-        model_manager.get_model_instance.return_value = mock_embedding_instance
-
-        mock_processor = MagicMock()
 
         # Mock ThreadPoolExecutor
         mock_future = MagicMock()
-        mock_future.result.return_value = 300  # Total tokens
+        mock_future.result.return_value = None
         mock_executor_instance = MagicMock()
         mock_executor_instance.__enter__.return_value = mock_executor_instance
         mock_executor_instance.__exit__.return_value = None
@@ -694,20 +689,51 @@ class TestIndexingRunnerLoad:
         mock_dependencies["executor"].return_value = mock_executor_instance
 
         # Mock update_document_index_status to avoid database calls
-        with patch.object(runner, "_update_document_index_status"):
+        with patch.object(runner, "_update_document_index_status") as mock_update_status:
             # Act
             runner._load(
-                mock_processor,
-                sample_dataset,
-                sample_dataset_document,
-                sample_documents,
-                mock_dependencies["session"],
+                session=mock_dependencies["session"],
+                dataset=sample_dataset,
+                dataset_document=sample_dataset_document,
+                documents=sample_documents,
+                total_tokens=300,
             )
 
         # Assert
-        model_manager.get_model_instance.assert_called_once()
+        mock_dependencies["model_manager"].assert_not_called()
         # Verify executor was used for parallel processing
         assert mock_executor_instance.submit.called
+        for submit_call in mock_executor_instance.submit.call_args_list:
+            assert submit_call.args[0] == runner._process_chunk
+            assert len(submit_call.args) == 6
+        mock_future.result.assert_called()
+        assert mock_update_status.call_args.kwargs["extra_update_params"][DatasetDocument.tokens] == 300
+
+    def test_load_propagates_worker_errors(
+        self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents
+    ):
+        runner = IndexingRunner()
+        mock_future = MagicMock()
+        mock_future.result.side_effect = RuntimeError("index failed")
+        mock_executor_instance = MagicMock()
+        mock_executor_instance.__enter__.return_value = mock_executor_instance
+        mock_executor_instance.__exit__.return_value = None
+        mock_executor_instance.submit.return_value = mock_future
+        mock_dependencies["executor"].return_value = mock_executor_instance
+
+        with (
+            patch.object(runner, "_update_document_index_status") as mock_update_status,
+            pytest.raises(RuntimeError, match="index failed"),
+        ):
+            runner._load(
+                session=mock_dependencies["session"],
+                dataset=sample_dataset,
+                dataset_document=sample_dataset_document,
+                documents=sample_documents,
+                total_tokens=300,
+            )
+
+        mock_update_status.assert_not_called()
 
     def test_load_with_economy_indexing(
         self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents
@@ -717,8 +743,6 @@ class TestIndexingRunnerLoad:
         runner = IndexingRunner()
         sample_dataset.indexing_technique = IndexTechniqueType.ECONOMY
 
-        mock_processor = MagicMock()
-
         # Mock thread for keyword indexing
         mock_thread_instance = MagicMock()
         mock_thread_instance.join = MagicMock()
@@ -728,11 +752,11 @@ class TestIndexingRunnerLoad:
         with patch.object(runner, "_update_document_index_status"):
             # Act
             runner._load(
-                mock_processor,
-                sample_dataset,
-                sample_dataset_document,
-                sample_documents,
-                mock_dependencies["session"],
+                session=mock_dependencies["session"],
+                dataset=sample_dataset,
+                dataset_document=sample_dataset_document,
+                documents=sample_documents,
+                total_tokens=0,
             )
 
         # Assert
@@ -759,16 +783,9 @@ class TestIndexingRunnerLoad:
                 )
             ]
 
-        mock_embedding_instance = MagicMock()
-        mock_embedding_instance.get_text_embedding_num_tokens.return_value = 50
-        model_manager = mock_dependencies["model_manager"].return_value
-        model_manager.get_model_instance.return_value = mock_embedding_instance
-
-        mock_processor = MagicMock()
-
         # Mock ThreadPoolExecutor
         mock_future = MagicMock()
-        mock_future.result.return_value = 150
+        mock_future.result.return_value = None
         mock_executor_instance = MagicMock()
         mock_executor_instance.__enter__.return_value = mock_executor_instance
         mock_executor_instance.__exit__.return_value = None
@@ -779,14 +796,15 @@ class TestIndexingRunnerLoad:
         with patch.object(runner, "_update_document_index_status"):
             # Act
             runner._load(
-                mock_processor,
-                sample_dataset,
-                sample_dataset_document,
-                sample_documents,
-                mock_dependencies["session"],
+                session=mock_dependencies["session"],
+                dataset=sample_dataset,
+                dataset_document=sample_dataset_document,
+                documents=sample_documents,
+                total_tokens=150,
             )
 
         # Assert
+        mock_dependencies["model_manager"].assert_not_called()
         # Verify no keyword thread for parent-child index
         mock_dependencies["thread"].assert_not_called()
 
@@ -850,6 +868,7 @@ class TestIndexingRunnerRun:
         segment.index_node_hash = "parent-hash"
         segment.document_id = dataset_document.id
         segment.dataset_id = dataset_document.dataset_id
+        segment.tokens = 12
         segment.get_child_chunks.return_value = [
             SimpleNamespace(content="child", index_node_id="child-node", index_node_hash="child-hash")
         ]
@@ -862,6 +881,32 @@ class TestIndexingRunnerRun:
 
         segment.get_child_chunks.assert_called_once_with(session=session)
         assert load.call_args.kwargs["documents"][0].children[0].page_content == "child"
+        assert load.call_args.kwargs["total_tokens"] == 12
+
+    def test_run_in_indexing_status_uses_tokens_from_all_segments(self, mock_dependencies, sample_dataset_documents):
+        runner = IndexingRunner()
+        dataset_document = sample_dataset_documents[0]
+        dataset = Mock(spec=Dataset)
+        completed_segment = Mock(spec=DocumentSegment)
+        completed_segment.status = SegmentStatus.COMPLETED
+        completed_segment.tokens = 10
+        incomplete_segment = Mock(spec=DocumentSegment)
+        incomplete_segment.status = SegmentStatus.WAITING
+        incomplete_segment.tokens = 20
+        incomplete_segment.content = "pending"
+        incomplete_segment.index_node_id = "pending-node"
+        incomplete_segment.index_node_hash = "pending-hash"
+        incomplete_segment.document_id = dataset_document.id
+        incomplete_segment.dataset_id = dataset_document.dataset_id
+        session = mock_dependencies["session"]
+        session.get.side_effect = lambda model, _: dataset_document if model is DatasetDocument else dataset
+        session.scalars.return_value.all.return_value = [completed_segment, incomplete_segment]
+
+        with patch.object(runner, "_load") as load:
+            runner.run_in_indexing_status(dataset_document, session)
+
+        assert load.call_args.kwargs["documents"][0].page_content == "pending"
+        assert load.call_args.kwargs["total_tokens"] == 30
 
     def test_run_success_single_document(self, mock_dependencies, sample_dataset_documents):
         """Test successful run with single document."""
@@ -953,6 +998,98 @@ class TestIndexingRunnerRun:
             with pytest.raises(DocumentIsPausedError):
                 runner.run([doc], mock_dependencies["session"])
 
+    def test_run_counts_each_transformed_document_once(self, mock_dependencies, sample_dataset_documents):
+        runner = IndexingRunner()
+        dataset_document = sample_dataset_documents[0]
+        dataset = Mock(spec=Dataset)
+        dataset.id = dataset_document.dataset_id
+        dataset.tenant_id = dataset_document.tenant_id
+        dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
+        current_user = Mock(spec=Account)
+        transformed_documents = [
+            Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}),
+            Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}),
+        ]
+        model_dispatch = {
+            DatasetDocument: dataset_document,
+            Dataset: dataset,
+            Account: current_user,
+        }
+        mock_dependencies["session"].get.side_effect = lambda model, _: model_dispatch.get(model)
+        process_rule = Mock(spec=DatasetProcessRule)
+        process_rule.to_dict.return_value = {"mode": "automatic", "rules": {}}
+        mock_dependencies["session"].scalar.return_value = process_rule
+
+        with (
+            patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]),
+            patch.object(runner, "_transform", return_value=transformed_documents),
+            patch.object(runner, "_load_segments") as load_segments,
+            patch.object(runner, "_load") as load,
+            patch(
+                "core.indexing_runner.calculate_segment_token_counts",
+                return_value=[11, 22],
+            ) as calculate_token_counts,
+        ):
+            runner.run([dataset_document], mock_dependencies["session"])
+
+        calculate_token_counts.assert_called_once_with(dataset=dataset, documents=transformed_documents)
+        load_segments.assert_called_once_with(
+            session=mock_dependencies["session"],
+            dataset=dataset,
+            dataset_document=dataset_document,
+            documents=transformed_documents,
+            token_counts=[11, 22],
+        )
+        assert load.call_args.kwargs["total_tokens"] == 33
+
+    def test_run_in_splitting_status_counts_each_transformed_document_once(
+        self, mock_dependencies, sample_dataset_documents
+    ):
+        runner = IndexingRunner()
+        dataset_document = sample_dataset_documents[0]
+        dataset_document.created_by = "user-1"
+        dataset = Mock(spec=Dataset)
+        dataset.id = dataset_document.dataset_id
+        dataset.tenant_id = dataset_document.tenant_id
+        dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY
+        current_user = Mock(spec=Account)
+        transformed_documents = [
+            Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}),
+            Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}),
+        ]
+        model_dispatch = {
+            DatasetDocument: dataset_document,
+            Dataset: dataset,
+            Account: current_user,
+        }
+        mock_dependencies["session"].get.side_effect = lambda model, _: model_dispatch.get(model)
+        mock_dependencies["session"].scalars.return_value.all.return_value = []
+        process_rule = Mock(spec=DatasetProcessRule)
+        process_rule.to_dict.return_value = {"mode": "automatic", "rules": {}}
+        mock_dependencies["session"].scalar.return_value = process_rule
+
+        with (
+            patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]),
+            patch.object(runner, "_transform", return_value=transformed_documents),
+            patch.object(runner, "_load_segments") as load_segments,
+            patch.object(runner, "_load") as load,
+            patch(
+                "core.indexing_runner.calculate_segment_token_counts",
+                return_value=[11, 22],
+            ) as calculate_token_counts,
+        ):
+            runner.run_in_splitting_status(dataset_document, mock_dependencies["session"])
+
+        calculate_token_counts.assert_called_once_with(dataset=dataset, documents=transformed_documents)
+        load_segments.assert_called_once_with(
+            session=mock_dependencies["session"],
+            dataset=dataset,
+            dataset_document=dataset_document,
+            documents=transformed_documents,
+            token_counts=[11, 22],
+        )
+        assert load.call_args.kwargs["total_tokens"] == 33
+
     def test_run_handles_provider_token_error(self, mock_dependencies, sample_dataset_documents):
         """Test run handles ProviderTokenNotInitError and updates document status."""
         # Arrange
@@ -1395,7 +1532,11 @@ class TestIndexingRunnerLoadSegments:
         ):
             # Act
             runner._load_segments(
-                sample_dataset, sample_dataset_document, sample_documents, mock_dependencies["session"]
+                session=mock_dependencies["session"],
+                dataset=sample_dataset,
+                dataset_document=sample_dataset_document,
+                documents=sample_documents,
+                token_counts=[10, 20],
             )
 
         # Assert
@@ -1405,7 +1546,10 @@ class TestIndexingRunnerLoadSegments:
             document_id=sample_dataset_document.id,
         )
         mock_docstore_instance.add_documents.assert_called_once_with(
-            docs=sample_documents, save_child=False, session=mock_dependencies["session"]
+            session=mock_dependencies["session"],
+            docs=sample_documents,
+            save_child=False,
+            token_counts=[10, 20],
         )
 
     def test_load_segments_parent_child_index(
@@ -1435,12 +1579,19 @@ class TestIndexingRunnerLoadSegments:
         ):
             # Act
             runner._load_segments(
-                sample_dataset, sample_dataset_document, sample_documents, mock_dependencies["session"]
+                session=mock_dependencies["session"],
+                dataset=sample_dataset,
+                dataset_document=sample_dataset_document,
+                documents=sample_documents,
+                token_counts=[10, 20],
             )
 
         # Assert
         mock_docstore_instance.add_documents.assert_called_once_with(
-            docs=sample_documents, save_child=True, session=mock_dependencies["session"]
+            session=mock_dependencies["session"],
+            docs=sample_documents,
+            save_child=True,
+            token_counts=[10, 20],
         )
 
     def test_load_segments_updates_word_count(
@@ -1462,7 +1613,11 @@ class TestIndexingRunnerLoadSegments:
         ):
             # Act
             runner._load_segments(
-                sample_dataset, sample_dataset_document, sample_documents, mock_dependencies["session"]
+                session=mock_dependencies["session"],
+                dataset=sample_dataset,
+                dataset_document=sample_dataset_document,
+                documents=sample_documents,
+                token_counts=[10, 20],
             )
 
         # Assert
@@ -1565,7 +1720,6 @@ class TestIndexingRunnerProcessChunk:
     """Unit tests for chunk processing in parallel.
 
     Tests cover:
-    - Token counting
     - Vector index creation
     - Segment status updates
     - Pause detection during processing
@@ -1590,16 +1744,12 @@ class TestIndexingRunnerProcessChunk:
         app.app_context.return_value.__exit__ = MagicMock()
         return app
 
-    def test_process_chunk_counts_tokens(self, mock_dependencies, mock_flask_app):
-        """Test process chunk correctly counts tokens."""
+    def test_process_chunk_loads_index_and_completes_segments(self, mock_dependencies, mock_flask_app):
+        """Test process chunk loads the index and completes segments without counting tokens."""
         # Arrange
         from core.indexing_runner import IndexingRunner
 
         runner = IndexingRunner()
-        mock_embedding_instance = MagicMock()
-        # Mock to return an iterable that sums to 150 tokens
-        mock_embedding_instance.get_text_embedding_num_tokens.return_value = [75, 75]
-
         mock_processor = MagicMock()
         chunk_documents = [
             Document(page_content="Chunk 1", metadata={"doc_id": "c1"}),
@@ -1638,18 +1788,19 @@ class TestIndexingRunnerProcessChunk:
             mock_factory.return_value.init_index_processor.return_value = mock_processor
 
             # Act - the method creates its own app_context and session
-            tokens = runner._process_chunk(
+            result = runner._process_chunk(
                 mock_flask_app,
                 IndexStructureType.PARAGRAPH_INDEX,
                 chunk_documents,
                 mock_dataset.id,
                 mock_dataset_document.id,
-                mock_embedding_instance,
             )
 
         # Assert
-        assert tokens == 150
+        assert result is None
         mock_processor.load.assert_called_once()
+        mock_dependencies["session"].execute.assert_called_once()
+        mock_dependencies["session"].commit.assert_called_once()
 
     def test_process_chunk_detects_pause(self, mock_dependencies, mock_flask_app):
         """Test process chunk detects document pause."""
@@ -1657,8 +1808,6 @@ class TestIndexingRunnerProcessChunk:
         from core.indexing_runner import IndexingRunner
 
         runner = IndexingRunner()
-        mock_embedding_instance = MagicMock()
-        mock_processor = MagicMock()
         chunk_documents = [Document(page_content="Chunk", metadata={"doc_id": "c1"})]
 
         mock_dataset = Mock(spec=Dataset)
@@ -1691,5 +1840,4 @@ class TestIndexingRunnerProcessChunk:
                     chunk_documents,
                     mock_dataset.id,
                     mock_dataset_document.id,
-                    mock_embedding_instance,
                 )

api/tests/unit_tests/core/tools/test_base_tool.py

@@ -18,6 +18,7 @@ from core.tools.entities.tool_entities import (
     ToolParameter,
     ToolProviderType,
 )
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID
 
 
 class DummyCastType:
@@ -365,6 +366,29 @@ def test_message_factory_helpers():
     assert json_message.message.json_object == {"k": "v"}
     assert json_message.message.suppress_output is True
 
+    ui_message = tool.create_ui_message(
+        {
+            "messages": [
+                {
+                    "version": "v0.9.1",
+                    "createSurface": {
+                        "surfaceId": "clock",
+                        "catalogId": A2UI_CATALOG_ID,
+                    },
+                },
+                {
+                    "version": "v0.9.1",
+                    "updateComponents": {
+                        "surfaceId": "clock",
+                        "components": [{"id": "root", "component": "Text", "text": "10:30"}],
+                    },
+                },
+            ]
+        }
+    )
+    assert ui_message.type == ToolInvokeMessage.MessageType.UI
+    assert ui_message.message.surface_id == "clock"
+
     variable_message = tool.create_variable_message("answer", 42, stream=False)
     assert variable_message.type == ToolInvokeMessage.MessageType.VARIABLE
     assert variable_message.message.variable_name == "answer"

api/tests/unit_tests/core/tools/test_builtin_tools_extra.py

@@ -2,9 +2,9 @@ from __future__ import annotations
 
 import calendar
 import math
-from datetime import date
+from datetime import UTC, date, datetime
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 from zoneinfo import ZoneInfo
 
 import pytest
@@ -28,6 +28,7 @@ from core.tools.builtin_tool.providers.webscraper.webscraper import WebscraperPr
 from core.tools.builtin_tool.tool import BuiltinTool
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID, ToolUIMessage
 from core.tools.errors import ToolInvokeError
 from graphon.file import FileType
 from graphon.model_runtime.entities.model_entities import ModelPropertyKey
@@ -53,10 +54,67 @@ def _raise_runtime_error(*_args: object, **_kwargs: object) -> None:
 
 def test_current_time_tool():
     current_tool = _build_builtin_tool(CurrentTimeTool)
-    utc_text = list(current_tool.invoke(session=MagicMock(), user_id="u", tool_parameters={"timezone": "UTC"}))[
-        0
-    ].message.text
-    assert utc_text
+    now = datetime(2024, 1, 1, 8, 30, tzinfo=UTC)
+
+    with patch(
+        "core.tools.builtin_tool.providers.time.tools.current_time.datetime",
+    ) as mock_datetime:
+        mock_datetime.now.return_value = now
+        messages = list(
+            current_tool.invoke(
+                session=MagicMock(),
+                user_id="u",
+                tool_parameters={"timezone": "UTC"},
+            )
+        )
+
+    assert [message.type for message in messages] == [
+        ToolInvokeMessage.MessageType.TEXT,
+        ToolInvokeMessage.MessageType.UI,
+    ]
+    assert messages[0].message.text == "2024-01-01 08:30:00 UTC"
+
+    ui_message = messages[1].message
+    assert isinstance(ui_message, ToolUIMessage)
+    assert ui_message.model_dump(mode="json", by_alias=True, exclude_none=True) == {
+        "protocol": "a2ui",
+        "protocol_version": "v0.9.1",
+        "messages": [
+            {
+                "version": "v0.9.1",
+                "createSurface": {
+                    "surfaceId": "current-time",
+                    "catalogId": A2UI_CATALOG_ID,
+                },
+            },
+            {
+                "version": "v0.9.1",
+                "updateDataModel": {
+                    "surfaceId": "current-time",
+                    "value": {"currentTime": "2024-01-01T08:30:00+00:00"},
+                },
+            },
+            {
+                "version": "v0.9.1",
+                "updateComponents": {
+                    "surfaceId": "current-time",
+                    "components": [
+                        {
+                            "id": "root",
+                            "component": "Card",
+                            "children": ["time"],
+                        },
+                        {
+                            "id": "time",
+                            "component": "DateTime",
+                            "value": {"path": "/currentTime"},
+                            "format": "datetime",
+                        },
+                    ],
+                },
+            },
+        ],
+    }
 
     invalid_tz = list(
         current_tool.invoke(session=MagicMock(), user_id="u", tool_parameters={"timezone": "Invalid/TZ"})

api/tests/unit_tests/core/tools/test_tool_engine.py

@@ -20,6 +20,7 @@ from core.tools.entities.tool_entities import (
     ToolParameter,
     ToolProviderType,
 )
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID, DIFY_UI_JSON_ENVELOPE_KEY
 from core.tools.errors import (
     ToolEngineInvokeError,
     ToolInvokeError,
@@ -210,7 +211,7 @@ def test_agent_invoke_success():
         ):
             with patch.object(ToolEngine, "_extract_tool_response_binary_and_text", return_value=iter([])):
                 with patch.object(ToolEngine, "_create_message_files", return_value=[]):
-                    result_text, message_files, result_meta = ToolEngine.agent_invoke(
+                    result = ToolEngine.agent_invoke(
                         session=MagicMock(),
                         tool=tool,
                         tool_parameters="hello",
@@ -221,20 +222,158 @@ def test_agent_invoke_success():
                         agent_tool_callback=callback,
                     )
 
-    assert result_text == "ok"
-    assert message_files == []
-    assert result_meta.error is None
+    assert result.observation == "ok"
+    assert result.message_files == []
+    assert result.ui_messages == []
+    assert result.meta.error is None
     callback.on_tool_start.assert_called_once()
     callback.on_tool_end.assert_called_once()
 
 
+def test_agent_invoke_extracts_ui_from_reserved_json_without_exposing_it_to_model():
+    tool = _build_tool(with_llm_parameter=True)
+    callback = Mock()
+    message = SimpleNamespace(id="m1", conversation_id="c1")
+    meta = ToolInvokeMeta.empty()
+    ui_payload = {
+        "protocol": "a2ui",
+        "protocol_version": "v0.9.1",
+        "messages": [
+            {
+                "version": "v0.9.1",
+                "createSurface": {"surfaceId": "clock", "catalogId": A2UI_CATALOG_ID},
+            },
+            {
+                "version": "v0.9.1",
+                "updateComponents": {
+                    "surfaceId": "clock",
+                    "components": [{"id": "root", "component": "Text", "text": "10:30"}],
+                },
+            },
+        ],
+    }
+    responses = [
+        tool.create_text_message("It is 10:30."),
+        tool.create_json_message({DIFY_UI_JSON_ENVELOPE_KEY: ui_payload}),
+        meta,
+    ]
+
+    with (
+        patch.object(ToolEngine, "_invoke", return_value=iter(responses)),
+        patch(
+            "core.tools.tool_engine.ToolFileMessageTransformer.transform_tool_invoke_messages",
+            side_effect=lambda messages, **kwargs: messages,
+        ),
+        patch.object(ToolEngine, "_extract_tool_response_binary_and_text", return_value=iter([])),
+        patch.object(ToolEngine, "_create_message_files", return_value=[]),
+    ):
+        result = ToolEngine.agent_invoke(
+            session=MagicMock(),
+            tool=tool,
+            tool_parameters={"query": "time"},
+            user_id="u1",
+            tenant_id="tenant-1",
+            message=message,
+            invoke_from=InvokeFrom.DEBUGGER,
+            agent_tool_callback=callback,
+        )
+
+    assert result.observation == "It is 10:30."
+    assert len(result.ui_messages) == 1
+    assert result.ui_messages[0].surface_id == "clock"
+    assert "createSurface" not in callback.on_tool_end.call_args.kwargs["tool_outputs"]
+
+
+def test_normalize_ui_messages_prefers_reserved_variable_channel():
+    tool = _build_tool()
+    ui_payload = {
+        "protocol": "a2ui",
+        "protocol_version": "v0.9.1",
+        "messages": [
+            {
+                "version": "v0.9.1",
+                "createSurface": {"surfaceId": "clock", "catalogId": A2UI_CATALOG_ID},
+            },
+            {
+                "version": "v0.9.1",
+                "updateComponents": {
+                    "surfaceId": "clock",
+                    "components": [{"id": "root", "component": "Text", "text": "10:30"}],
+                },
+            },
+        ],
+    }
+    variable = tool.create_variable_message(DIFY_UI_JSON_ENVELOPE_KEY, ui_payload)
+
+    normalized = list(ToolEngine.normalize_ui_messages([variable]))
+
+    assert len(normalized) == 1
+    assert normalized[0].type == ToolInvokeMessage.MessageType.UI
+    assert normalized[0].message.surface_id == "clock"
+    assert ToolEngine.tool_response_to_str(normalized) == ""
+
+
+def test_collect_agent_messages_drops_oversized_ui_batch_but_preserves_observation():
+    tool = _build_tool()
+    ui_messages = [
+        tool.create_ui_message(
+            {
+                "messages": [
+                    {
+                        "version": "v0.9.1",
+                        "createSurface": {
+                            "surfaceId": f"surface-{index}",
+                            "catalogId": A2UI_CATALOG_ID,
+                        },
+                    },
+                    {
+                        "version": "v0.9.1",
+                        "updateComponents": {
+                            "surfaceId": f"surface-{index}",
+                            "components": [{"id": "root", "component": "Text", "text": str(index)}],
+                        },
+                    },
+                ]
+            }
+        )
+        for index in range(17)
+    ]
+
+    messages, extracted_ui = ToolEngine.collect_agent_messages(
+        [
+            tool.create_text_message("weather observation before;"),
+            *ui_messages,
+            tool.create_text_message("weather observation after"),
+        ]
+    )
+
+    assert extracted_ui == []
+    assert all(message.type != ToolInvokeMessage.MessageType.UI for message in messages)
+    assert ToolEngine.tool_response_to_str(messages) == "weather observation before;weather observation after"
+
+
+def test_collect_agent_messages_drops_malformed_reserved_ui_but_preserves_observation():
+    tool = _build_tool()
+    malformed_ui = tool.create_variable_message(
+        DIFY_UI_JSON_ENVELOPE_KEY,
+        {"protocol": "a2ui", "messages": []},
+    )
+
+    messages, extracted_ui = ToolEngine.collect_agent_messages(
+        [tool.create_text_message("weather observation"), malformed_ui]
+    )
+
+    assert extracted_ui == []
+    assert ToolEngine.tool_response_to_str(messages) == "weather observation"
+
+
 def test_agent_invoke_param_validation_error():
     tool = _build_tool(with_llm_parameter=True)
     callback = Mock()
     message = SimpleNamespace(id="m1", conversation_id="c1")
 
     with patch.object(ToolEngine, "_invoke", side_effect=ToolParameterValidationError("bad-param")):
-        error_text, files, error_meta = ToolEngine.agent_invoke(
+        result = ToolEngine.agent_invoke(
             session=MagicMock(),
             tool=tool,
             tool_parameters={"a": 1},
@@ -245,9 +384,9 @@ def test_agent_invoke_param_validation_error():
             agent_tool_callback=callback,
         )
 
-    assert "tool parameters validation error" in error_text
-    assert files == []
-    assert error_meta.error
+    assert "tool parameters validation error" in result.observation
+    assert result.message_files == []
+    assert result.meta.error
 
 
 def test_agent_invoke_engine_meta_error():
@@ -257,7 +396,7 @@ def test_agent_invoke_engine_meta_error():
     engine_error = ToolEngineInvokeError(ToolInvokeMeta.error_instance("meta failure"))
 
     with patch.object(ToolEngine, "_invoke", side_effect=engine_error):
-        error_text, files, error_meta = ToolEngine.agent_invoke(
+        result = ToolEngine.agent_invoke(
             session=MagicMock(),
             tool=tool,
             tool_parameters={"a": 1},
@@ -268,9 +407,9 @@ def test_agent_invoke_engine_meta_error():
             agent_tool_callback=callback,
         )
 
-    assert "meta failure" in error_text
-    assert files == []
-    assert error_meta.error == "meta failure"
+    assert "meta failure" in result.observation
+    assert result.message_files == []
+    assert result.meta.error == "meta failure"
 
 
 def test_convert_tool_response_excludes_variable_messages():
@@ -301,7 +440,7 @@ def test_agent_invoke_tool_invoke_error():
     message = SimpleNamespace(id="m1", conversation_id="c1")
 
     with patch.object(ToolEngine, "_invoke", side_effect=ToolInvokeError("invoke boom")):
-        error_text, files, _ = ToolEngine.agent_invoke(
+        result = ToolEngine.agent_invoke(
             session=MagicMock(),
             tool=tool,
             tool_parameters={"a": 1},
@@ -312,5 +451,5 @@ def test_agent_invoke_tool_invoke_error():
             agent_tool_callback=callback,
         )
 
-    assert "tool invoke error" in error_text
-    assert files == []
+    assert "tool invoke error" in result.observation
+    assert result.message_files == []

api/tests/unit_tests/core/tools/test_ui_entities.py

@@ -0,0 +1,500 @@
+from __future__ import annotations
+
+import json
+import math
+from copy import deepcopy
+
+import pytest
+from pydantic import ValidationError
+
+from core.tools.entities.ui_entities import (
+    A2UI_CATALOG_ID,
+    DIFY_UI_JSON_ENVELOPE_KEY,
+    MAX_UI_PARTS_PAYLOAD_BYTES,
+    MAX_UI_PARTS_PER_MESSAGE,
+    A2UIComponent,
+    MessageUIPart,
+    ToolUIMessage,
+    build_ui_part_id,
+    extract_ui_message_from_json,
+    parse_history_ui_parts,
+    parse_tool_ui_messages,
+    validate_tool_ui_message_batch,
+    validate_ui_part_batch,
+)
+
+
+def _valid_ui_message() -> dict:
+    return {
+        "protocol": "a2ui",
+        "protocol_version": "v0.9.1",
+        "messages": [
+            {
+                "version": "v0.9.1",
+                "createSurface": {
+                    "surfaceId": "weather",
+                    "catalogId": A2UI_CATALOG_ID,
+                },
+            },
+            {
+                "version": "v0.9.1",
+                "updateDataModel": {
+                    "surfaceId": "weather",
+                    "value": {"temperature": 23},
+                },
+            },
+            {
+                "version": "v0.9.1",
+                "updateComponents": {
+                    "surfaceId": "weather",
+                    "components": [
+                        {
+                            "id": "root",
+                            "component": "Card",
+                            "title": "Shanghai",
+                            "children": ["temperature"],
+                        },
+                        {
+                            "id": "temperature",
+                            "component": "Metric",
+                            "label": "Temperature",
+                            "value": {"path": "/temperature"},
+                            "unit": "°C",
+                        },
+                    ],
+                },
+            },
+        ],
+        "fallback": "Shanghai: 23°C",
+    }
+
+
+def _ui_message(surface_id: str, *, large: bool = False) -> ToolUIMessage:
+    value = deepcopy(_valid_ui_message())
+    for message in value["messages"]:
+        operation = message.get("createSurface") or message.get("updateDataModel") or message.get("updateComponents")
+        operation["surfaceId"] = surface_id
+    if large:
+        value["messages"][1]["updateDataModel"]["value"] = ["x" * 4096] * 20
+    return ToolUIMessage.model_validate(value)
+
+
+def _ui_part(index: int, *, large: bool = False, sequence: int = 1) -> MessageUIPart:
+    ui_message = _ui_message(f"surface-{index}", large=large)
+    return MessageUIPart.from_tool_ui_message(
+        part_id=f"call-{index}:{ui_message.surface_id}",
+        sequence=sequence,
+        ui_message=ui_message,
+    )
+
+
+def test_tool_ui_message_accepts_fixed_catalog_surface() -> None:
+    ui_message = ToolUIMessage.model_validate(_valid_ui_message())
+
+    assert ui_message.surface_id == "weather"
+    assert ui_message.messages[-1].update_components is not None
+    part = MessageUIPart.from_tool_ui_message(
+        part_id="call-1:weather",
+        sequence=1,
+        ui_message=ui_message,
+    )
+    assert part.part_id == "call-1:weather"
+    assert part.messages[0].create_surface.catalog_id == A2UI_CATALOG_ID  # type: ignore[union-attr]
+
+
+def test_tool_ui_message_rejects_custom_component_props() -> None:
+    invalid = _valid_ui_message()
+    invalid["messages"][-1]["updateComponents"]["components"][1]["className"] = "custom"
+
+    with pytest.raises(ValidationError):
+        ToolUIMessage.model_validate(invalid)
+
+
+def test_tool_ui_message_rejects_image_and_omits_it_from_the_catalog_schema() -> None:
+    invalid = _valid_ui_message()
+    invalid["messages"][-1]["updateComponents"]["components"] = [
+        {
+            "id": "root",
+            "component": "Image",
+            "src": "https://attacker.example/tracker.png",
+            "alt": "tracker",
+        }
+    ]
+
+    with pytest.raises(ValidationError):
+        ToolUIMessage.model_validate(invalid)
+
+    schema = A2UIComponent.model_json_schema()
+    assert "Image" not in json.dumps(schema)
+    assert "src" not in schema["properties"]
+    assert "alt" not in schema["properties"]
+
+
+def test_data_model_allows_non_executable_arbitrary_keys() -> None:
+    value = _valid_ui_message()
+    value["messages"][1]["updateDataModel"]["value"] = {
+        "action": "forecast",
+        "theme": "weather",
+        "style": "windy",
+    }
+
+    ui_message = ToolUIMessage.model_validate(value)
+
+    assert ui_message.messages[1].update_data_model is not None
+    assert ui_message.messages[1].update_data_model.value["action"] == "forecast"
+
+
+@pytest.mark.parametrize("unsafe_key", ["__proto__", "constructor", "prototype"])
+def test_data_model_rejects_prototype_pollution_keys(unsafe_key: str) -> None:
+    value = _valid_ui_message()
+    value["messages"][1]["updateDataModel"]["value"] = {"safe": {unsafe_key: "polluted"}}
+
+    with pytest.raises(ValidationError, match="forbidden keys"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_data_model_rejects_excessive_depth_and_nodes() -> None:
+    too_deep: object = "leaf"
+    for _ in range(17):
+        too_deep = {"nested": too_deep}
+    value = _valid_ui_message()
+    value["messages"][1]["updateDataModel"]["value"] = too_deep
+    with pytest.raises(ValidationError, match="nested too deeply"):
+        ToolUIMessage.model_validate(value)
+
+    value = _valid_ui_message()
+    value["messages"][1]["updateDataModel"]["value"] = list(range(2000))
+    with pytest.raises(ValidationError, match="too many nodes"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_data_model_rejects_cumulative_depth_from_pointer_patches() -> None:
+    value = _valid_ui_message()
+    value["messages"][1] = {
+        "version": "v0.9.1",
+        "updateDataModel": {
+            "surfaceId": "weather",
+            "path": "/" + "/".join(f"level-{index}" for index in range(16)),
+            "value": {"leaf": "too deep after materialization"},
+        },
+    }
+
+    with pytest.raises(ValidationError, match="nested too deeply"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_data_model_rejects_cumulative_nodes_from_object_patches() -> None:
+    value = _valid_ui_message()
+    value["messages"][1:1] = [
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "value": {"left": list(range(999))},
+            },
+        },
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "path": "/right",
+                "value": list(range(999)),
+            },
+        },
+    ]
+    value["messages"].pop(3)
+
+    with pytest.raises(ValidationError, match="too many nodes"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_data_model_materializes_root_replacements_and_array_patches() -> None:
+    value = _valid_ui_message()
+    value["messages"][1:2] = [
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "value": {"discarded": list(range(1500))},
+            },
+        },
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "path": "/",
+                "value": {"days": []},
+            },
+        },
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "path": "/days/0",
+                "value": {"temperature": 23},
+            },
+        },
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "path": "/days/1",
+                "value": {"temperature": 24},
+            },
+        },
+    ]
+
+    ui_message = ToolUIMessage.model_validate(value)
+
+    assert len(ui_message.messages) == 6
+
+
+def test_data_model_rejects_array_patch_gaps() -> None:
+    value = _valid_ui_message()
+    value["messages"][1:2] = [
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "value": {"days": []},
+            },
+        },
+        {
+            "version": "v0.9.1",
+            "updateDataModel": {
+                "surfaceId": "weather",
+                "path": "/days/2",
+                "value": {"temperature": 23},
+            },
+        },
+    ]
+
+    with pytest.raises(ValidationError, match="array index"):
+        ToolUIMessage.model_validate(value)
+
+
+@pytest.mark.parametrize(
+    ("component", "props"),
+    [
+        ("Metric", {"label": "Temperature", "value": math.inf}),
+        ("KeyValue", {"label": "Humidity", "value": math.nan}),
+        ("Progress", {"label": "Loading", "value": -math.inf}),
+    ],
+)
+def test_components_reject_non_finite_numbers(component: str, props: dict[str, object]) -> None:
+    value = _valid_ui_message()
+    value["messages"][-1]["updateComponents"]["components"] = [{"id": "root", "component": component, **props}]
+
+    with pytest.raises(ValidationError):
+        ToolUIMessage.model_validate(value)
+
+
+@pytest.mark.parametrize("number", [math.inf, -math.inf, math.nan])
+def test_data_model_rejects_non_finite_numbers(number: float) -> None:
+    value = _valid_ui_message()
+    value["messages"][1]["updateDataModel"]["value"] = {"number": number}
+
+    with pytest.raises(ValidationError):
+        ToolUIMessage.model_validate(value)
+
+
+def test_progress_requires_an_accessible_label() -> None:
+    value = _valid_ui_message()
+    value["messages"][-1]["updateComponents"]["components"] = [{"id": "root", "component": "Progress", "value": 50}]
+
+    with pytest.raises(ValidationError, match="requires properties.*label"):
+        ToolUIMessage.model_validate(value)
+
+
+@pytest.mark.parametrize(
+    ("mutate", "expected"),
+    [
+        (
+            lambda value: value["messages"].pop(0),
+            "first A2UI message must be createSurface",
+        ),
+        (
+            lambda value: value["messages"][-1]["updateComponents"]["components"].pop(0),
+            "component with id 'root'",
+        ),
+        (
+            lambda value: value["messages"][-1]["updateComponents"]["components"][0].update({"children": ["missing"]}),
+            "unknown child",
+        ),
+    ],
+)
+def test_tool_ui_message_rejects_invalid_surface_graph(mutate, expected: str) -> None:
+    invalid = _valid_ui_message()
+    mutate(invalid)
+
+    with pytest.raises(ValidationError, match=expected):
+        ToolUIMessage.model_validate(invalid)
+
+
+def test_tool_ui_message_rejects_component_cycle() -> None:
+    invalid = _valid_ui_message()
+    invalid["messages"][-1]["updateComponents"]["components"][1] = {
+        "id": "temperature",
+        "component": "Column",
+        "children": ["root"],
+    }
+
+    with pytest.raises(ValidationError, match="cycle"):
+        ToolUIMessage.model_validate(invalid)
+
+
+def test_tool_ui_message_accepts_rooted_component_tree() -> None:
+    value = _valid_ui_message()
+    value["messages"][-1]["updateComponents"]["components"] = [
+        {"id": "root", "component": "Column", "children": ["summary", "details"]},
+        {"id": "summary", "component": "Row", "children": ["temperature", "condition"]},
+        {
+            "id": "temperature",
+            "component": "Metric",
+            "label": "Temperature",
+            "value": 23,
+            "unit": "°C",
+        },
+        {"id": "condition", "component": "Badge", "text": "Sunny"},
+        {"id": "details", "component": "Text", "text": "Light wind"},
+    ]
+
+    ui_message = ToolUIMessage.model_validate(value)
+
+    assert ui_message.messages[-1].update_components is not None
+    assert len(ui_message.messages[-1].update_components.components) == 5
+
+
+def test_tool_ui_message_rejects_duplicate_children() -> None:
+    value = _valid_ui_message()
+    value["messages"][-1]["updateComponents"]["components"][0]["children"] = [
+        "temperature",
+        "temperature",
+    ]
+
+    with pytest.raises(ValidationError, match="duplicate ids"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_tool_ui_message_rejects_root_as_child() -> None:
+    value = _valid_ui_message()
+    value["messages"][-1]["updateComponents"]["components"].append(
+        {"id": "orphan", "component": "Column", "children": ["root"]}
+    )
+
+    with pytest.raises(ValidationError, match="root component cannot be referenced"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_tool_ui_message_rejects_component_with_multiple_parents() -> None:
+    value = _valid_ui_message()
+    value["messages"][-1]["updateComponents"]["components"] = [
+        {"id": "root", "component": "Row", "children": ["left", "right"]},
+        {"id": "left", "component": "Column", "children": ["shared"]},
+        {"id": "right", "component": "Column", "children": ["shared"]},
+        {"id": "shared", "component": "Text", "text": "Shared"},
+    ]
+
+    with pytest.raises(ValidationError, match="multiple parents"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_tool_ui_message_rejects_unreachable_component() -> None:
+    value = _valid_ui_message()
+    value["messages"][-1]["updateComponents"]["components"].append(
+        {"id": "orphan", "component": "Text", "text": "Not reachable"}
+    )
+
+    with pytest.raises(ValidationError, match="unreachable from root"):
+        ToolUIMessage.model_validate(value)
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "/weather/~2value",
+        "/weather/__proto__/polluted",
+        "/constructor/value",
+    ],
+)
+def test_tool_ui_message_rejects_unsafe_json_pointers(path: str) -> None:
+    invalid_binding = _valid_ui_message()
+    invalid_binding["messages"][-1]["updateComponents"]["components"][1]["value"] = {"path": path}
+    with pytest.raises(ValidationError):
+        ToolUIMessage.model_validate(invalid_binding)
+
+    invalid_update = _valid_ui_message()
+    invalid_update["messages"][1]["updateDataModel"]["path"] = path
+    with pytest.raises(ValidationError):
+        ToolUIMessage.model_validate(invalid_update)
+
+
+def test_data_model_update_path_rejects_large_array_index() -> None:
+    value = _valid_ui_message()
+    value["messages"][1]["updateDataModel"]["path"] = "/items/1001"
+
+    with pytest.raises(ValidationError, match="array index"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_json_pointer_rejects_more_than_sixteen_segments() -> None:
+    value = _valid_ui_message()
+    value["messages"][1]["updateDataModel"]["path"] = "/" + "/".join(f"level-{index}" for index in range(17))
+
+    with pytest.raises(ValidationError, match="too many segments"):
+        ToolUIMessage.model_validate(value)
+
+
+def test_extract_ui_message_from_legacy_json_envelope() -> None:
+    parsed = extract_ui_message_from_json({DIFY_UI_JSON_ENVELOPE_KEY: _valid_ui_message()})
+
+    assert parsed is not None
+    assert parsed.surface_id == "weather"
+    assert extract_ui_message_from_json({"normal": "json"}) is None
+
+
+def test_tool_ui_batch_enforces_count_and_payload_limits() -> None:
+    validate_tool_ui_message_batch([_ui_message(f"surface-{index}") for index in range(16)])
+
+    with pytest.raises(ValueError, match="more than"):
+        validate_tool_ui_message_batch([_ui_message(f"surface-{index}") for index in range(17)])
+
+    large_messages = [_ui_message(f"large-{index}", large=True) for index in range(7)]
+    with pytest.raises(ValueError, match=str(MAX_UI_PARTS_PAYLOAD_BYTES)):
+        validate_tool_ui_message_batch(large_messages)
+
+
+def test_parse_tool_ui_messages_rejects_count_before_parsing_items() -> None:
+    with pytest.raises(ValueError, match="more than"):
+        parse_tool_ui_messages([object() for _ in range(MAX_UI_PARTS_PER_MESSAGE + 1)])
+
+
+def test_ui_part_id_is_stable_unambiguous_and_bounded() -> None:
+    first = build_ui_part_id("call:a", "b")
+    second = build_ui_part_id("call", "a:b")
+
+    assert first == build_ui_part_id("call:a", "b")
+    assert first != second
+    assert len(build_ui_part_id("x" * 10_000, "weather")) <= 512
+    assert len(build_ui_part_id("\ud800", "weather")) <= 512
+
+
+def test_history_ui_parts_ignore_invalid_stale_and_over_budget_items() -> None:
+    first = _ui_part(0)
+    newer = first.model_copy(update={"sequence": 2})
+    raw_parts = [
+        first.model_dump(mode="json"),
+        {"part_id": "invalid"},
+        newer.model_dump(mode="json"),
+        *[_ui_part(index).model_dump(mode="json") for index in range(1, 18)],
+    ]
+
+    parts = parse_history_ui_parts(raw_parts)
+
+    assert len(parts) == MAX_UI_PARTS_PER_MESSAGE
+    assert parts[0].sequence == 2
+    validate_ui_part_batch(parts)
+
+    large_parts = parse_history_ui_parts([_ui_part(index, large=True).model_dump(mode="json") for index in range(7)])
+    assert len(large_parts) < 7
+    validate_ui_part_batch(large_parts)

api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_runtime.py

@@ -13,6 +13,7 @@ from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCal
 from core.plugin.impl.exc import PluginDaemonClientSideError, PluginInvokeError
 from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.entities.tool_entities import ToolProviderType as CoreToolProviderType
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID
 from core.tools.errors import ToolInvokeError
 from core.tools.tool_engine import ToolEngine
 from core.tools.tool_manager import ToolManager
@@ -327,6 +328,29 @@ def test_convert_message_payload_rejects_unknown_types(runtime: DifyToolNodeRunt
         runtime._convert_message_payload(object())
 
 
+def test_adapt_messages_ignores_chat_ui_messages(runtime: DifyToolNodeRuntime) -> None:
+    ui_message = ToolInvokeMessage(
+        type=ToolInvokeMessage.MessageType.UI,
+        message={
+            "messages": [
+                {
+                    "version": "v0.9.1",
+                    "createSurface": {"surfaceId": "clock", "catalogId": A2UI_CATALOG_ID},
+                },
+                {
+                    "version": "v0.9.1",
+                    "updateComponents": {
+                        "surfaceId": "clock",
+                        "components": [{"id": "root", "component": "Text", "text": "10:30"}],
+                    },
+                },
+            ]
+        },
+    )
+
+    assert list(runtime._adapt_messages(iter([ui_message]), provider_name="provider")) == []
+
+
 def test_resolve_provider_icons_prefers_builtin_tool_icons(runtime: DifyToolNodeRuntime) -> None:
     plugin = SimpleNamespace(
         plugin_id="langgenius/tools",

api/tests/unit_tests/extensions/otel/test_runtime.py

@@ -0,0 +1,75 @@
+from unittest import mock
+
+import pytest
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+
+from core.logging.context import clear_request_context
+
+
+@pytest.fixture(autouse=True)
+def _reset_logging_context():
+    clear_request_context()
+    yield
+    clear_request_context()
+
+
+def test_on_user_loaded_does_not_write_to_non_recording_span() -> None:
+    from extensions.otel import runtime
+
+    span = mock.MagicMock()
+    span.is_recording.return_value = False
+    user = mock.Mock(id="user-id")
+
+    with (
+        mock.patch.object(runtime.dify_config, "ENABLE_OTEL", True),
+        mock.patch("opentelemetry.trace.get_current_span", return_value=span),
+        mock.patch.object(runtime, "extract_tenant_id", return_value="tenant-id"),
+    ):
+        runtime.on_user_loaded(None, user)
+
+    span.is_recording.assert_called_once_with()
+    span.set_attribute.assert_not_called()
+    span.set_attributes.assert_not_called()
+
+
+def test_on_user_loaded_sets_attributes_on_recording_span() -> None:
+    from extensions.otel import runtime
+    from extensions.otel.semconv import DifySpanAttributes, GenAIAttributes
+
+    span = mock.MagicMock()
+    span.is_recording.return_value = True
+    user = mock.Mock(id="user-id")
+
+    with (
+        mock.patch.object(runtime.dify_config, "ENABLE_OTEL", True),
+        mock.patch("opentelemetry.trace.get_current_span", return_value=span),
+        mock.patch.object(runtime, "extract_tenant_id", return_value="tenant-id"),
+    ):
+        runtime.on_user_loaded(None, user)
+
+    span.set_attributes.assert_called_once_with(
+        {
+            DifySpanAttributes.TENANT_ID: "tenant-id",
+            GenAIAttributes.USER_ID: "user-id",
+        }
+    )
+
+
+def test_on_user_loaded_ignores_ended_sdk_span(caplog) -> None:
+    from extensions.otel import runtime
+
+    tracer_provider = TracerProvider()
+    span = tracer_provider.get_tracer(__name__).start_span("ended")
+    span.end()
+    user = mock.Mock(id="user-id")
+
+    with (
+        trace.use_span(span, end_on_exit=False),
+        mock.patch.object(runtime.dify_config, "ENABLE_OTEL", True),
+        mock.patch.object(runtime, "extract_tenant_id", return_value="tenant-id"),
+        caplog.at_level("WARNING", logger="opentelemetry.sdk.trace"),
+    ):
+        runtime.on_user_loaded(None, user)
+
+    assert "Setting attribute on ended span" not in caplog.text

api/tests/unit_tests/extensions/test_ext_login.py

@@ -1,10 +1,22 @@
 import json
+from typing import cast
+from unittest import mock
 
+import pytest
 from flask import Response
 
+from core.logging.context import clear_request_context, get_identity_context
+from extensions import ext_login
 from extensions.ext_login import unauthorized_handler
 
 
+@pytest.fixture(autouse=True)
+def _reset_logging_context():
+    clear_request_context()
+    yield
+    clear_request_context()
+
+
 def test_unauthorized_handler_returns_json_response() -> None:
     response = unauthorized_handler()
 
@@ -15,3 +27,50 @@ def test_unauthorized_handler_returns_json_response() -> None:
         "code": "unauthorized",
         "message": "Unauthorized.",
     }
+
+
+def test_on_user_logged_in_sets_account_logging_identity() -> None:
+    account = mock.Mock(spec=ext_login.Account)
+    account.id = "account-id"
+    account.current_tenant_id = "tenant-id"
+    clear_request_context()
+
+    ext_login.on_user_logged_in(None, account)
+
+    assert get_identity_context() == ("tenant-id", "account-id", "account")
+
+
+def test_on_user_logged_in_sets_end_user_logging_identity() -> None:
+    end_user = mock.Mock(spec=ext_login.EndUser)
+    end_user.id = "end-user-id"
+    end_user.tenant_id = "tenant-id"
+    end_user.type = "browser"
+    clear_request_context()
+
+    ext_login.on_user_logged_in(None, end_user)
+
+    assert get_identity_context() == ("tenant-id", "end-user-id", "browser")
+
+
+def test_on_user_logged_in_does_not_break_auth_when_identity_is_unavailable() -> None:
+    account = mock.Mock(spec=ext_login.Account)
+    type(account).current_tenant_id = mock.PropertyMock(side_effect=RuntimeError("unavailable"))
+    account.id = "account-id"
+    clear_request_context()
+
+    with mock.patch.object(ext_login.logger, "exception") as logger_exception:
+        ext_login.on_user_logged_in(None, account)
+
+    assert get_identity_context() == ("", "", "")
+    logger_exception.assert_called_once_with("Failed to set logging identity context")
+
+
+def test_on_user_logged_in_logs_unsupported_user_type() -> None:
+    unsupported_user = cast(ext_login.LoginUser, object())
+    clear_request_context()
+
+    with mock.patch.object(ext_login.logger, "exception") as logger_exception:
+        ext_login.on_user_logged_in(None, unsupported_user)
+
+    assert get_identity_context() == ("", "", "")
+    logger_exception.assert_called_once_with("Failed to set logging identity context")

api/tests/unit_tests/fields/test_message_fields.py

@@ -1,5 +1,6 @@
 from decimal import Decimal
 
+from core.tools.entities.ui_entities import A2UI_CATALOG_ID, MessageUIPart, ToolUIMessage
 from fields.message_fields import ExploreMessageListItem, MessageListItem, WebMessageListItem
 
 
@@ -18,6 +19,44 @@ def _base_kwargs():
     }
 
 
+def _ui_part_payload(index: int, *, large: bool = False) -> dict:
+    surface_id = f"surface-{index}"
+    messages = [
+        {
+            "version": "v0.9.1",
+            "createSurface": {
+                "surfaceId": surface_id,
+                "catalogId": A2UI_CATALOG_ID,
+            },
+        },
+    ]
+    if large:
+        messages.append(
+            {
+                "version": "v0.9.1",
+                "updateDataModel": {
+                    "surfaceId": surface_id,
+                    "value": ["x" * 4096] * 20,
+                },
+            }
+        )
+    messages.append(
+        {
+            "version": "v0.9.1",
+            "updateComponents": {
+                "surfaceId": surface_id,
+                "components": [{"id": "root", "component": "Text", "text": "Weather"}],
+            },
+        }
+    )
+    ui_message = ToolUIMessage(messages=messages)
+    return MessageUIPart.from_tool_ui_message(
+        part_id=f"call-{index}:{surface_id}",
+        sequence=1,
+        ui_message=ui_message,
+    ).model_dump(mode="json")
+
+
 class TestExploreMessageListItem:
     def test_exposes_metadata_for_history_rehydration(self):
         # The Explore/installed-app surface must surface message_metadata (incl. reasoning)
@@ -66,3 +105,46 @@ class TestExploreMessageListItem:
         assert payload["message_tokens"] == 7
         assert payload["answer_tokens"] == 11
         assert payload["total_tokens"] == 18
+
+    def test_service_message_exposes_only_valid_bounded_ui_parts(self):
+        metadata = {
+            "usage": {"total_tokens": 18},
+            "ui_parts": [
+                _ui_part_payload(0),
+                {"part_id": "broken"},
+                *[_ui_part_payload(index) for index in range(1, 18)],
+            ],
+        }
+
+        payload = MessageListItem(
+            **_base_kwargs(),
+            message_metadata_dict=metadata,
+        ).model_dump(mode="json")
+
+        assert "metadata" not in payload
+        assert len(payload["ui_parts"]) == 16
+        assert payload["ui_parts"][0]["part_id"] == "call-0:surface-0"
+
+    def test_web_metadata_stays_raw_while_top_level_ui_parts_are_safe(self):
+        metadata = {
+            "reasoning": {"llm": "thinking..."},
+            "ui_parts": [_ui_part_payload(0), {"part_id": "broken"}],
+        }
+
+        payload = WebMessageListItem(
+            **_base_kwargs(),
+            message_metadata_dict=metadata,
+        ).model_dump(mode="json")
+
+        assert payload["metadata"] == metadata
+        assert len(payload["ui_parts"]) == 1
+
+    def test_service_message_ui_parts_respect_cumulative_payload_budget(self):
+        metadata = {"ui_parts": [_ui_part_payload(index, large=True) for index in range(7)]}
+
+        payload = MessageListItem(
+            **_base_kwargs(),
+            message_metadata_dict=metadata,
+        ).model_dump(mode="json")
+
+        assert 0 < len(payload["ui_parts"]) < 7

api/tests/unit_tests/services/agent/test_agent_observability_service.py

@@ -5,7 +5,8 @@ from types import SimpleNamespace
 import pytest
 
 from core.app.entities.app_invoke_entities import InvokeFrom
-from models.enums import ConversationFromSource, MessageStatus
+from graphon.enums import WorkflowNodeExecutionStatus
+from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus
 from services.agent import observability_service as observability_service_module
 from services.agent.observability_service import AgentLogQueryParams, AgentObservabilityService
 
@@ -268,6 +269,183 @@ def test_list_logs_sorts_by_requested_field(monkeypatch: pytest.MonkeyPatch) ->
     assert [item["id"] for item in payload["data"]] == ["old", "new"]
 
 
+def test_list_log_messages_merges_deduplicates_and_sorts_sources(monkeypatch: pytest.MonkeyPatch) -> None:
+    service = AgentObservabilityService(session=None)
+    webapp_message = SimpleNamespace(id="shared")
+    webapp_row = {"id": "shared", "created_at": 10, "updated_at": 30}
+    workflow_rows = [
+        {"id": "shared", "created_at": 10, "updated_at": 20},
+        {"id": "workflow-only", "created_at": 20, "updated_at": 10},
+    ]
+    monkeypatch.setattr(service, "_list_webapp_messages", lambda **kwargs: [webapp_message])
+    monkeypatch.setattr(service, "serialize_log_message", lambda message: webapp_row)
+    monkeypatch.setattr(service, "_list_workflow_messages", lambda **kwargs: workflow_rows)
+
+    payload = service.list_log_messages(
+        app=SimpleNamespace(id="agent-app"),  # type: ignore[arg-type]
+        agent_id="agent-1",
+        conversation_id="execution-1",
+        params=AgentLogQueryParams(
+            sources=("webapp:agent-app", "workflow:workflow-app"),
+            sort_by="created_at",
+            sort_order="asc",
+        ),
+    )
+
+    assert payload == {
+        "data": [workflow_rows[0], workflow_rows[1]],
+        "page": 1,
+        "limit": 20,
+        "total": 2,
+        "has_more": False,
+    }
+
+
+def test_list_workflow_logs_uses_node_executions_without_messages() -> None:
+    created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC)
+    node_execution = SimpleNamespace(
+        id="node-execution-1",
+        title="Agent",
+        status=WorkflowNodeExecutionStatus.SUCCEEDED,
+        created_by_role=CreatorUserRole.END_USER,
+        created_by="end-user-1",
+        created_at=created_at,
+        finished_at=None,
+    )
+    workflow_app = SimpleNamespace(
+        id="workflow-app-1",
+        name="Marketing Department",
+        icon_type=None,
+        icon=None,
+        icon_background=None,
+    )
+
+    class FakeRow:
+        node_execution_id = node_execution.id
+        node_title = node_execution.title
+        node_status = node_execution.status
+        node_created_by_role = node_execution.created_by_role
+        node_created_by = node_execution.created_by
+        node_created_at = node_execution.created_at
+        node_finished_at = node_execution.finished_at
+        workflow_id = "workflow-1"
+        workflow_version = "v1"
+        node_id = "node-1"
+
+        def __getitem__(self, index: int):
+            return (None, None, None, None, None, None, None, workflow_app)[index]
+
+    class FakeResult:
+        def all(self):
+            return [FakeRow()]
+
+    class FakeSession:
+        def __init__(self):
+            self.query = ""
+
+        def execute(self, stmt):
+            self.query = str(stmt)
+            return FakeResult()
+
+    session = FakeSession()
+    service = AgentObservabilityService(session)
+
+    rows = service._list_workflow_conversation_logs(
+        app=SimpleNamespace(tenant_id="tenant-1"),  # type: ignore[arg-type]
+        agent_id="agent-1",
+        params=AgentLogQueryParams(),
+        source_filter=AgentObservabilityService.resolve_source_filter("workflow:workflow-app-1"),
+    )
+
+    assert "FROM workflow_node_executions" in session.query
+    assert "JOIN messages" not in session.query
+    assert rows[0]["id"] == "node-execution-1"
+    assert rows[0]["source"]["app_name"] == "Marketing Department"
+
+
+def test_list_workflow_messages_uses_node_execution_identity(monkeypatch: pytest.MonkeyPatch) -> None:
+    node_execution = SimpleNamespace(id="node-execution-1")
+
+    class FakeScalarResult:
+        def all(self):
+            return [node_execution]
+
+    class FakeSession:
+        def __init__(self):
+            self.query = ""
+
+        def scalars(self, stmt):
+            self.query = str(stmt)
+            return FakeScalarResult()
+
+    session = FakeSession()
+    service = AgentObservabilityService(session)
+    serialized = {"id": "node-execution-1", "conversation_id": "node-execution-1"}
+    monkeypatch.setattr(service, "serialize_workflow_node_message", lambda execution: serialized)
+
+    rows = service._list_workflow_messages(
+        app=SimpleNamespace(tenant_id="tenant-1"),  # type: ignore[arg-type]
+        agent_id="agent-1",
+        conversation_id="node-execution-1",
+        params=AgentLogQueryParams(),
+        source_filter=AgentObservabilityService.resolve_source_filter("workflow:workflow-app-1"),
+    )
+
+    assert rows == [serialized]
+    assert "FROM workflow_node_executions" in session.query
+    assert "workflow_node_executions.id =" in session.query
+    assert "JOIN messages" not in session.query
+
+
+def test_apply_workflow_node_filters_supports_time_keyword_and_status() -> None:
+    class FakeStmt:
+        def __init__(self):
+            self.conditions = []
+
+        def where(self, *conditions):
+            self.conditions.extend(conditions)
+            return self
+
+    stmt = FakeStmt()
+    params = AgentLogQueryParams(
+        start=datetime(2026, 7, 1, tzinfo=UTC),
+        end=datetime(2026, 8, 1, tzinfo=UTC),
+        keyword="meeting_100%",
+        statuses=("success",),
+    )
+
+    result = AgentObservabilityService._apply_workflow_node_filters(
+        stmt,
+        params=params,
+        workflow_app=SimpleNamespace(name=observability_service_module.App.name),
+    )
+
+    assert result is stmt
+    assert len(stmt.conditions) == 4
+
+
+def test_apply_workflow_node_status_filter_supports_all_status_groups() -> None:
+    class FakeStmt:
+        def __init__(self):
+            self.conditions = []
+
+        def where(self, *conditions):
+            self.conditions.extend(conditions)
+            return self
+
+    stmt = FakeStmt()
+
+    result = AgentObservabilityService._apply_workflow_node_status_filter(stmt, ("normal", "error", "paused"))
+
+    assert result is stmt
+    assert len(stmt.conditions) == 1
+    empty_stmt = FakeStmt()
+    assert AgentObservabilityService._apply_workflow_node_status_filter(empty_stmt, ()) is empty_stmt
+    assert empty_stmt.conditions == []
+    with pytest.raises(ValueError, match="Unsupported status"):
+        AgentObservabilityService._apply_workflow_node_status_filter(FakeStmt(), ("unknown",))
+
+
 def test_source_serializers_return_structured_frontend_shape() -> None:
     app = SimpleNamespace(
         id="app-1",
@@ -400,6 +578,152 @@ def test_serialize_log_message_returns_frontend_log_shape() -> None:
     }
 
 
+def test_serialize_workflow_node_message_returns_frontend_log_shape() -> None:
+    created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC)
+    finished_at = datetime(2026, 7, 23, 7, 0, 28, tzinfo=UTC)
+    node_execution = SimpleNamespace(
+        id="node-execution-1",
+        title="Agent",
+        inputs=(
+            '{"agent_backend_request":{"composition":{"layers":['
+            '{"name":"workflow_node_job_prompt","config":{"user":"Summarize the meeting"}},'
+            '{"name":"workflow_user_prompt","config":{"user":"Focus on action items"}}]}}}'
+        ),
+        outputs='{"output":"Alice owns the follow-up."}',
+        execution_metadata=(
+            '{"agent_log":{"agent_backend":{"usage":{"prompt_tokens":10,"completion_tokens":5,'
+            '"total_tokens":15,"total_price":"0.0015","currency":"USD","latency":1.25}}}}'
+        ),
+        status=WorkflowNodeExecutionStatus.SUCCEEDED,
+        error=None,
+        elapsed_time=1.5,
+        created_by_role=CreatorUserRole.END_USER,
+        created_by="end-user-1",
+        created_at=created_at,
+        finished_at=finished_at,
+    )
+
+    payload = AgentObservabilityService.serialize_workflow_node_message(node_execution)  # type: ignore[arg-type]
+
+    assert payload == {
+        "id": "node-execution-1",
+        "message_id": "node-execution-1",
+        "conversation_id": "node-execution-1",
+        "query": "Summarize the meeting\n\nFocus on action items",
+        "answer": "Alice owns the follow-up.",
+        "status": "success",
+        "error": None,
+        "from_end_user_id": "end-user-1",
+        "from_account_id": None,
+        "message_tokens": 10,
+        "answer_tokens": 5,
+        "total_tokens": 15,
+        "total_price": "0.0015",
+        "currency": "USD",
+        "latency": 1.25,
+        "created_at": int(created_at.timestamp()),
+        "updated_at": int(finished_at.timestamp()),
+    }
+
+
+def test_serialize_workflow_node_message_handles_sparse_runtime_data() -> None:
+    created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC)
+    node_execution = SimpleNamespace(
+        id="node-execution-2",
+        title="Fallback prompt",
+        inputs={
+            "agent_backend_request": {
+                "composition": {
+                    "layers": [
+                        None,
+                        {"name": "unrelated", "config": {"user": "ignored"}},
+                        {"name": "workflow_user_prompt", "config": {"user": "  "}},
+                    ]
+                }
+            }
+        },
+        outputs={"output": {"structured": True}},
+        execution_metadata={
+            "agent_log": {
+                "agent_backend": {
+                    "usage": {
+                        "prompt_tokens": "2",
+                        "completion_tokens": 3,
+                    }
+                }
+            }
+        },
+        status=WorkflowNodeExecutionStatus.PAUSED,
+        error=None,
+        elapsed_time=2,
+        created_by_role=CreatorUserRole.ACCOUNT.value,
+        created_by="account-1",
+        created_at=created_at,
+        finished_at=None,
+    )
+
+    payload = AgentObservabilityService.serialize_workflow_node_message(node_execution)  # type: ignore[arg-type]
+
+    assert payload["query"] == "Fallback prompt"
+    assert payload["answer"] == '{"output": {"structured": true}}'
+    assert payload["status"] == "paused"
+    assert payload["from_end_user_id"] is None
+    assert payload["from_account_id"] == "account-1"
+    assert payload["total_tokens"] == 5
+    assert payload["total_price"] == "0"
+    assert payload["currency"] == ""
+    assert payload["latency"] == 2.0
+    assert payload["updated_at"] == int(created_at.timestamp())
+
+
+def test_workflow_node_serialization_helpers_handle_invalid_values() -> None:
+    assert AgentObservabilityService._json_mapping(None) == {}
+    assert AgentObservabilityService._json_mapping("not-json") == {}
+    assert AgentObservabilityService._json_mapping("[]") == {}
+    assert AgentObservabilityService._mapping_value({"value": []}, "value") == {}
+    assert AgentObservabilityService._int_value(None) == 0
+    assert AgentObservabilityService._int_value("not-a-number") == 0
+    assert (
+        AgentObservabilityService._workflow_node_query(
+            {"agent_backend_request": {"composition": {"layers": "invalid"}}}, fallback="fallback"
+        )
+        == "fallback"
+    )
+    assert AgentObservabilityService._workflow_node_answer({"output": 1, "text": "fallback text"}) == "fallback text"
+    assert AgentObservabilityService._workflow_node_answer({}) == ""
+
+
+def test_serialize_workflow_execution_log_uses_node_execution_identity() -> None:
+    created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC)
+    node_execution = SimpleNamespace(
+        id="node-execution-1",
+        title="Agent",
+        status=WorkflowNodeExecutionStatus.FAILED,
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by="account-1",
+        created_at=created_at,
+        finished_at=None,
+    )
+
+    payload = AgentObservabilityService._serialize_workflow_execution_log(
+        node_execution_id=node_execution.id,
+        title=node_execution.title,
+        status=node_execution.status,
+        created_by_role=node_execution.created_by_role,
+        created_by=node_execution.created_by,
+        created_at=node_execution.created_at,
+        finished_at=node_execution.finished_at,
+        source={"id": "workflow:app-1:workflow-1:v1:node-1"},
+    )
+
+    assert payload["id"] == "node-execution-1"
+    assert payload["conversation_id"] == "node-execution-1"
+    assert payload["message_count"] == 1
+    assert payload["end_user_id"] is None
+    assert payload["status"] == "failed"
+    assert payload["unread"] is False
+
+
 def test_build_charts_and_summary_match_monitoring_metrics() -> None:
     rows = [
         {

api/tests/unit_tests/services/agent/test_agent_services.py

@@ -3496,8 +3496,61 @@ class TestAgentAppBackingAgent:
         assert session.deleted == []
         assert session.commits == 1
 
-    def test_refresh_agent_app_debug_conversation_enqueues_cleanup_for_old_runtime_sessions(
+    def test_refresh_agent_app_debug_conversation_rotates_preview_mapping_each_time(
         self, monkeypatch: pytest.MonkeyPatch
+    ):
+        agent = Agent(
+            id="agent-1",
+            tenant_id="tenant-1",
+            name="Iris",
+            description="",
+            agent_kind=AgentKind.DIFY_AGENT,
+            scope=AgentScope.ROSTER,
+            source=AgentSource.AGENT_APP,
+            status=AgentStatus.ACTIVE,
+            app_id="app-1",
+        )
+        session = FakeSession(scalar=[agent, None])
+        service = AgentRosterService(session)
+        cleanup = MagicMock()
+        monkeypatch.setattr(service, "_cleanup_debug_conversation_runtime_sessions", cleanup)
+
+        first_conversation_id = service.refresh_agent_app_debug_conversation_id(
+            tenant_id="tenant-1",
+            agent_id="agent-1",
+            account_id="account-1",
+            draft_type=AgentConfigDraftType.DRAFT,
+        )
+        mapping = next(value for value in session.added if isinstance(value, AgentDebugConversation))
+        session._scalar.extend([agent, mapping])
+        second_conversation_id = service.refresh_agent_app_debug_conversation_id(
+            tenant_id="tenant-1",
+            agent_id="agent-1",
+            account_id="account-1",
+            draft_type=AgentConfigDraftType.DRAFT,
+        )
+
+        assert first_conversation_id != second_conversation_id
+        assert mapping.draft_type == AgentConfigDraftType.DRAFT
+        assert mapping.conversation_id == second_conversation_id
+        assert len([value for value in session.added if isinstance(value, Conversation)]) == 2
+        cleanup.assert_called_once_with(
+            tenant_id="tenant-1",
+            agent_id="agent-1",
+            account_id="account-1",
+            draft_type=AgentConfigDraftType.DRAFT,
+            app_id="app-1",
+            conversation_id=first_conversation_id,
+        )
+
+    @pytest.mark.parametrize(
+        "draft_type",
+        [AgentConfigDraftType.DRAFT, AgentConfigDraftType.DEBUG_BUILD],
+    )
+    def test_refresh_agent_app_debug_conversation_enqueues_cleanup_for_old_runtime_sessions(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        draft_type: AgentConfigDraftType,
     ):
         agent = Agent(
             id="agent-1",
@@ -3525,9 +3578,21 @@ class TestAgentAppBackingAgent:
         )
         session = FakeSession(scalar=[agent, mapping])
         service = AgentRosterService(session)
+        events: list[str] = []
+        original_commit = session.commit
+
+        def record_commit() -> None:
+            events.append("commit")
+            original_commit()
+
+        def list_active_sessions(**kwargs: object) -> list[object]:
+            events.append("cleanup")
+            return [stored_session]
+
         cleanup_delay = MagicMock()
         cleanup_store = MagicMock()
-        cleanup_store.list_active_sessions_for_conversation.return_value = [stored_session]
+        cleanup_store.list_active_sessions_for_conversation.side_effect = list_active_sessions
+        monkeypatch.setattr(session, "commit", record_commit)
         monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store)
         monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay)
 
@@ -3535,6 +3600,7 @@ class TestAgentAppBackingAgent:
             tenant_id="tenant-1",
             agent_id="agent-1",
             account_id="account-1",
+            draft_type=draft_type,
         )
 
         cleanup_store.list_active_sessions_for_conversation.assert_called_once_with(
@@ -3546,10 +3612,10 @@ class TestAgentAppBackingAgent:
         payload = cleanup_delay.call_args.args[0]
         assert payload["metadata"]["conversation_id"] == "old-conversation"
         assert payload["metadata"]["agent_id"] == "agent-9"
-        assert payload["metadata"]["draft_type"] == "debug_build"
+        assert payload["metadata"]["draft_type"] == draft_type.value
         assert (
             payload["idempotency_key"]
-            == "tenant-1:agent-1:account-1:debug_build:old-conversation:debug-session-cleanup:"
+            == f"tenant-1:agent-1:account-1:{draft_type.value}:old-conversation:debug-session-cleanup:"
             "agent-9:snap-9:run-old"
         )
         cleanup_store.mark_cleaned.assert_called_once_with(
@@ -3558,6 +3624,41 @@ class TestAgentAppBackingAgent:
         )
         assert mapping.app_id == "app-1"
         assert mapping.conversation_id == conversation_id
+        assert events == ["commit", "cleanup"]
+
+    def test_refresh_agent_app_debug_conversation_does_not_cleanup_when_commit_fails(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        agent = Agent(
+            id="agent-1",
+            tenant_id="tenant-1",
+            name="Iris",
+            description="",
+            agent_kind=AgentKind.DIFY_AGENT,
+            scope=AgentScope.ROSTER,
+            source=AgentSource.AGENT_APP,
+            status=AgentStatus.ACTIVE,
+            app_id="app-1",
+        )
+        mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation")
+        session = FakeSession(scalar=[agent, mapping])
+        service = AgentRosterService(session)
+        runtime_store_factory = MagicMock()
+        cleanup_delay = MagicMock()
+        monkeypatch.setattr(session, "commit", MagicMock(side_effect=RuntimeError("database unavailable")))
+        monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", runtime_store_factory)
+        monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay)
+
+        with pytest.raises(RuntimeError, match="database unavailable"):
+            service.refresh_agent_app_debug_conversation_id(
+                tenant_id="tenant-1",
+                agent_id="agent-1",
+                account_id="account-1",
+                draft_type=AgentConfigDraftType.DRAFT,
+            )
+
+        runtime_store_factory.assert_not_called()
+        cleanup_delay.assert_not_called()
 
     def test_refresh_agent_app_debug_conversation_marks_old_runtime_sessions_clean_when_enqueue_fails(
         self, monkeypatch: pytest.MonkeyPatch

api/tests/unit_tests/services/plugin/test_plugin_service_installation.py

@@ -24,7 +24,6 @@ from models.engine import db
 from models.provider import Provider, ProviderCredential, TenantPreferredModelProvider
 from services.errors.plugin import PluginInstallationForbiddenError
 from services.feature_service import (
-    DeploymentEdition,
     PluginInstallationPermissionModel,
     PluginInstallationScope,
     SystemFeatureModel,
@@ -36,11 +35,10 @@ def _make_features(
     scope: PluginInstallationScope = PluginInstallationScope.ALL,
 ) -> SystemFeatureModel:
     return SystemFeatureModel(
-        deployment_edition=DeploymentEdition.COMMUNITY,
         plugin_installation_permission=PluginInstallationPermissionModel(
             restrict_to_marketplace_only=restrict_to_marketplace,
             plugin_installation_scope=scope,
-        ),
+        )
     )
 
 

api/tests/unit_tests/services/test_agent_tool_inner_service.py

@@ -7,6 +7,7 @@ import pytest
 from sqlalchemy.orm import Session
 
 from core.tools.entities.tool_entities import ToolInvokeMessage, ToolProviderType
+from core.tools.entities.ui_entities import ToolUIMessage
 from core.tools.errors import (
     ToolInvokeError,
     ToolParameterValidationError,
@@ -79,6 +80,74 @@ def _messages() -> Generator[ToolInvokeMessage, None, None]:
     )
 
 
+def _ui_messages() -> Generator[ToolInvokeMessage, None, None]:
+    yield ToolInvokeMessage(
+        type=ToolInvokeMessage.MessageType.TEXT,
+        message=ToolInvokeMessage.TextMessage(text="Sunny"),
+    )
+    yield ToolInvokeMessage(
+        type=ToolInvokeMessage.MessageType.UI,
+        message=ToolUIMessage.model_validate(
+            {
+                "protocol": "a2ui",
+                "protocol_version": "v0.9.1",
+                "messages": [
+                    {
+                        "version": "v0.9.1",
+                        "createSurface": {
+                            "surfaceId": "weather",
+                            "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                        },
+                    },
+                    {
+                        "version": "v0.9.1",
+                        "updateComponents": {
+                            "surfaceId": "weather",
+                            "components": [
+                                {
+                                    "id": "root",
+                                    "component": "Text",
+                                    "text": "Sunny",
+                                }
+                            ],
+                        },
+                    },
+                ],
+            }
+        ),
+    )
+
+
+def _oversized_ui_messages() -> Generator[ToolInvokeMessage, None, None]:
+    yield ToolInvokeMessage(
+        type=ToolInvokeMessage.MessageType.TEXT,
+        message=ToolInvokeMessage.TextMessage(text="Sunny"),
+    )
+    for index in range(17):
+        surface_id = f"weather-{index}"
+        yield ToolInvokeMessage(
+            type=ToolInvokeMessage.MessageType.UI,
+            message=ToolUIMessage(
+                messages=[
+                    {
+                        "version": "v0.9.1",
+                        "createSurface": {
+                            "surfaceId": surface_id,
+                            "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                        },
+                    },
+                    {
+                        "version": "v0.9.1",
+                        "updateComponents": {
+                            "surfaceId": surface_id,
+                            "components": [{"id": "root", "component": "Text", "text": "Sunny"}],
+                        },
+                    },
+                ]
+            ),
+        )
+
+
 @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
 def test_invoke_uses_agent_tool_runtime_and_returns_observation(sqlite_session: Session) -> None:
     fake_tool = MagicMock()
@@ -111,6 +180,74 @@ def test_invoke_uses_agent_tool_runtime_and_returns_observation(sqlite_session:
     assert sqlite_session.in_transaction()
 
 
+@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
+def test_invoke_keeps_ui_in_messages_and_out_of_observation(sqlite_session: Session) -> None:
+    _persist_app(sqlite_session)
+
+    with (
+        patch("services.agent_tool_inner_service.ToolManager.get_agent_tool_runtime", return_value=MagicMock()),
+        patch("services.agent_tool_inner_service.ToolEngine.generic_invoke", return_value=_ui_messages()),
+        patch(
+            "services.agent_tool_inner_service.ToolFileMessageTransformer.transform_tool_invoke_messages",
+            side_effect=lambda messages, **_kwargs: messages,
+        ),
+    ):
+        response = AgentToolInnerService().invoke(_request(), session=sqlite_session)
+
+    assert response.observation == "Sunny"
+    assert response.messages[1] == {
+        "type": "ui",
+        "message": {
+            "protocol": "a2ui",
+            "protocol_version": "v0.9.1",
+            "messages": [
+                {
+                    "version": "v0.9.1",
+                    "createSurface": {
+                        "surfaceId": "weather",
+                        "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                    },
+                },
+                {
+                    "version": "v0.9.1",
+                    "updateComponents": {
+                        "surfaceId": "weather",
+                        "components": [
+                            {
+                                "id": "root",
+                                "component": "Text",
+                                "text": "Sunny",
+                            }
+                        ],
+                    },
+                },
+            ],
+        },
+        "meta": None,
+    }
+
+
+@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
+def test_invoke_drops_oversized_ui_batch_without_losing_observation(sqlite_session: Session) -> None:
+    _persist_app(sqlite_session)
+
+    with (
+        patch("services.agent_tool_inner_service.ToolManager.get_agent_tool_runtime", return_value=MagicMock()),
+        patch(
+            "services.agent_tool_inner_service.ToolEngine.generic_invoke",
+            return_value=_oversized_ui_messages(),
+        ),
+        patch(
+            "services.agent_tool_inner_service.ToolFileMessageTransformer.transform_tool_invoke_messages",
+            side_effect=lambda messages, **_kwargs: messages,
+        ),
+    ):
+        response = AgentToolInnerService().invoke(_request(), session=sqlite_session)
+
+    assert response.observation == "Sunny"
+    assert [message["type"] for message in response.messages] == ["text"]
+
+
 @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
 def test_invoke_raises_app_not_found_when_session_has_no_app(sqlite_session: Session) -> None:
     with pytest.raises(AgentToolInnerServiceError) as exc_info:

api/tests/unit_tests/services/test_feature_service_change_email.py

@@ -1,25 +0,0 @@
-import pytest
-
-from services.feature_service import FeatureService
-
-
-@pytest.mark.parametrize("enabled", [False, True])
-def test_get_system_features_reads_enable_change_email(
-    monkeypatch: pytest.MonkeyPatch,
-    enabled: bool,
-) -> None:
-    monkeypatch.setattr("services.feature_service.dify_config.ENABLE_CHANGE_EMAIL", enabled)
-
-    result = FeatureService.get_system_features()
-
-    assert result.enable_change_email is enabled
-
-
-def test_enterprise_disables_change_email(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr("services.feature_service.dify_config.ENABLE_CHANGE_EMAIL", True)
-    monkeypatch.setattr("services.feature_service.dify_config.ENTERPRISE_ENABLED", True)
-    monkeypatch.setattr("services.feature_service.FeatureService._fulfill_params_from_enterprise", lambda *_: None)
-
-    result = FeatureService.get_system_features()
-
-    assert result.enable_change_email is False

api/tests/unit_tests/services/test_feature_service_deployment_edition.py

@@ -1,34 +0,0 @@
-import pytest
-from pydantic import ValidationError
-
-from services.feature_service import DeploymentEdition, FeatureService, SystemFeatureModel
-
-
-def test_system_feature_model_requires_deployment_edition() -> None:
-    with pytest.raises(ValidationError):
-        SystemFeatureModel.model_validate({})
-
-
-@pytest.mark.parametrize(
-    ("edition", "enterprise_enabled", "expected"),
-    [
-        ("SELF_HOSTED", False, DeploymentEdition.COMMUNITY),
-        ("SELF_HOSTED", True, DeploymentEdition.ENTERPRISE),
-        ("CLOUD", False, DeploymentEdition.CLOUD),
-        ("CLOUD", True, DeploymentEdition.CLOUD),
-    ],
-)
-def test_get_system_features_resolves_deployment_edition(
-    monkeypatch: pytest.MonkeyPatch,
-    edition: str,
-    enterprise_enabled: bool,
-    expected: DeploymentEdition,
-) -> None:
-    monkeypatch.setattr("services.feature_service.dify_config.EDITION", edition)
-    monkeypatch.setattr("services.feature_service.dify_config.ENTERPRISE_ENABLED", enterprise_enabled)
-    monkeypatch.setattr("services.feature_service.FeatureService._fulfill_params_from_enterprise", lambda *_: None)
-
-    result = FeatureService.get_system_features()
-
-    assert result.deployment_edition is expected
-    assert result.model_dump(mode="json")["deployment_edition"] == expected.value

api/tests/unit_tests/services/test_feature_service_enable_app_deploy.py

@@ -1,7 +1,7 @@
 import pytest
 
 from services import feature_service as feature_service_module
-from services.feature_service import DeploymentEdition, FeatureService, SystemFeatureModel
+from services.feature_service import FeatureService, SystemFeatureModel
 
 
 @pytest.mark.parametrize(
@@ -29,7 +29,7 @@ def test_fulfill_params_from_enterprise_enable_app_deploy(
         staticmethod(lambda: enterprise_info),
     )
 
-    features = SystemFeatureModel(deployment_edition=DeploymentEdition.COMMUNITY)
+    features = SystemFeatureModel()
     features.enable_app_deploy = initial
 
     FeatureService._fulfill_params_from_enterprise(features)

api/tests/unit_tests/services/test_feature_service_knowledge_fs.py

@@ -1,10 +1,10 @@
 import pytest
 
-from services.feature_service import DeploymentEdition, FeatureService, SystemFeatureModel
+from services.feature_service import FeatureService, SystemFeatureModel
 
 
 def test_system_feature_model_disables_knowledge_fs_by_default() -> None:
-    assert SystemFeatureModel(deployment_edition=DeploymentEdition.COMMUNITY).knowledge_fs_enabled is False
+    assert SystemFeatureModel().knowledge_fs_enabled is False
 
 
 @pytest.mark.parametrize("enabled", [False, True])

api/tests/unit_tests/services/test_feature_service_learn_app.py

@@ -1,14 +1,12 @@
 import pytest
 
 from services import feature_service as feature_service_module
-from services.feature_service import DeploymentEdition, FeatureService, SystemFeatureModel
+from services.feature_service import FeatureService, SystemFeatureModel
 
 
 def test_system_feature_model_defaults_enable_learn_app():
-    system_features = SystemFeatureModel(deployment_edition=DeploymentEdition.COMMUNITY)
-
-    assert system_features.enable_learn_app is True
-    assert system_features.enable_step_by_step_tour is False
+    assert SystemFeatureModel().enable_learn_app is True
+    assert SystemFeatureModel().enable_step_by_step_tour is False
 
 
 @pytest.mark.parametrize("enabled", [True, False])

api/tests/unit_tests/services/test_feature_service_licensed_seats.py

@@ -1,7 +1,7 @@
 import pytest
 
 from services import feature_service as feature_service_module
-from services.feature_service import DeploymentEdition, FeatureService, SystemFeatureModel
+from services.feature_service import FeatureService, SystemFeatureModel
 
 _ENTERPRISE_INFO = {"License": {"licensedSeats": {"enabled": True, "limit": 3, "used": 1}}}
 
@@ -14,7 +14,7 @@ def test_fulfill_params_from_enterprise_parses_licensed_seats(monkeypatch: pytes
         staticmethod(lambda: _ENTERPRISE_INFO),
     )
 
-    features = SystemFeatureModel(deployment_edition=DeploymentEdition.COMMUNITY)
+    features = SystemFeatureModel()
     FeatureService._fulfill_params_from_enterprise(features, is_authenticated=True)
 
     assert features.license.seats.enabled is True
@@ -30,7 +30,7 @@ def test_fulfill_params_from_enterprise_withholds_seats_when_unauthenticated(mon
         staticmethod(lambda: _ENTERPRISE_INFO),
     )
 
-    features = SystemFeatureModel(deployment_edition=DeploymentEdition.COMMUNITY)
+    features = SystemFeatureModel()
     FeatureService._fulfill_params_from_enterprise(features, is_authenticated=False)
 
     assert features.license.seats.enabled is False

api/tests/unit_tests/services/test_feature_service_webapp_public_access.py

@@ -1,6 +1,6 @@
 import pytest
 
-from services.feature_service import DeploymentEdition, FeatureService, SystemFeatureModel
+from services.feature_service import FeatureService, SystemFeatureModel
 
 
 @pytest.mark.parametrize(
@@ -18,7 +18,7 @@ def test_fulfill_system_params_from_env_sets_allow_public_access(
 ):
     monkeypatch.setattr("services.feature_service.dify_config.WEBAPP_PUBLIC_ACCESS_ENABLED", env_value)
 
-    system_features = SystemFeatureModel(deployment_edition=DeploymentEdition.COMMUNITY)
+    system_features = SystemFeatureModel()
     FeatureService._fulfill_system_params_from_env(system_features)
 
     assert system_features.webapp_auth.allow_public_access is expected

api/tests/unit_tests/services/test_recommended_app_service.py

@@ -12,7 +12,7 @@ from sqlalchemy.orm import Session
 
 from models.model import AccountTrialAppRecord, App, AppMode, TrialApp
 from services import recommended_app_service as service_module
-from services.feature_service import DeploymentEdition, SystemFeatureModel
+from services.feature_service import SystemFeatureModel
 from services.recommended_app_service import RecommendedAppService
 
 pytestmark = pytest.mark.parametrize(
@@ -298,10 +298,7 @@ class TestRecommendedAppServiceGetDetail:
         sqlite_session: Session,
     ) -> None:
         mock_config.HOSTED_FETCH_APP_TEMPLATES_MODE = "remote"
-        mock_feature_service.get_system_features.return_value = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_trial_app=False,
-        )
+        mock_feature_service.get_system_features.return_value = SystemFeatureModel(enable_trial_app=False)
         cases: list[tuple[str, RecommendedAppPayload]] = [
             (
                 "complex-app",
@@ -340,10 +337,7 @@ class TestRecommendedAppServiceGetDetail:
         mock_feature_service: MagicMock,
         sqlite_session: Session,
     ) -> None:
-        mock_feature_service.get_system_features.return_value = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_trial_app=False,
-        )
+        mock_feature_service.get_system_features.return_value = SystemFeatureModel(enable_trial_app=False)
         for mode in ["remote", "builtin", "db"]:
             mock_config.HOSTED_FETCH_APP_TEMPLATES_MODE = mode
             detail = _app_detail(app_id="test-app", name=f"App from {mode}")
@@ -373,10 +367,7 @@ class TestRecommendedAppServiceGetLearnDifyApps:
         sqlite_session: Session,
     ) -> None:
         mock_config.HOSTED_FETCH_APP_TEMPLATES_MODE = "remote"
-        mock_feature_service.get_system_features.return_value = SystemFeatureModel(
-            deployment_edition=DeploymentEdition.COMMUNITY,
-            enable_trial_app=False,
-        )
+        mock_feature_service.get_system_features.return_value = SystemFeatureModel(enable_trial_app=False)
         expected_app = RecommendedAppPayload(app_id="app-1", category="Workflow")
         mock_instance = MagicMock()
         mock_instance.get_learn_dify_apps.return_value = {
@@ -411,12 +402,7 @@ class TestRecommendedAppServiceGetLearnDifyApps:
         monkeypatch.setattr(
             service_module.FeatureService,
             "get_system_features",
-            MagicMock(
-                return_value=SystemFeatureModel(
-                    deployment_edition=DeploymentEdition.COMMUNITY,
-                    enable_trial_app=True,
-                )
-            ),
+            MagicMock(return_value=SystemFeatureModel(enable_trial_app=True)),
         )
         can_trial_mock = MagicMock(return_value=True)
         monkeypatch.setattr(RecommendedAppService, "_can_trial_app", can_trial_mock)
@@ -439,12 +425,7 @@ class TestRecommendedAppServiceTrialFeatures:
         monkeypatch.setattr(
             service_module.FeatureService,
             "get_system_features",
-            MagicMock(
-                return_value=SystemFeatureModel(
-                    deployment_edition=DeploymentEdition.COMMUNITY,
-                    enable_trial_app=False,
-                )
-            ),
+            MagicMock(return_value=SystemFeatureModel(enable_trial_app=False)),
         )
 
         result = RecommendedAppService.get_recommended_apps_and_categories("en-US", session=sqlite_session)
@@ -475,12 +456,7 @@ class TestRecommendedAppServiceTrialFeatures:
         monkeypatch.setattr(
             service_module.FeatureService,
             "get_system_features",
-            MagicMock(
-                return_value=SystemFeatureModel(
-                    deployment_edition=DeploymentEdition.COMMUNITY,
-                    enable_trial_app=True,
-                )
-            ),
+            MagicMock(return_value=SystemFeatureModel(enable_trial_app=True)),
         )
 
         result = RecommendedAppService.get_recommended_apps_and_categories("ja-JP", session=sqlite_session)
@@ -516,12 +492,7 @@ class TestRecommendedAppServiceTrialFeatures:
         monkeypatch.setattr(
             service_module.FeatureService,
             "get_system_features",
-            MagicMock(
-                return_value=SystemFeatureModel(
-                    deployment_edition=DeploymentEdition.COMMUNITY,
-                    enable_trial_app=True,
-                )
-            ),
+            MagicMock(return_value=SystemFeatureModel(enable_trial_app=True)),
         )
 
         result = RecommendedAppService.get_recommend_app_detail(app_id, session=sqlite_session)

dify-agent/src/dify_agent/layers/dify_core_tools/layer.py

@@ -9,12 +9,14 @@ provider-local state stay in the API process.
 
 from __future__ import annotations
 
+import json
+from collections.abc import Mapping, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
 from typing import ClassVar
 
 import httpx
-from pydantic_ai import RunContext, Tool
+from pydantic_ai import RunContext, Tool, ToolReturn
 from pydantic_ai.tools import ToolDefinition
 from typing_extensions import Self, override
 
@@ -34,6 +36,10 @@ from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer
 
 CORE_TOOL_STRICT = False
 TEMPORARY_UNAVAILABLE_OBSERVATION = "Tool is temporarily unavailable. Please continue without it if possible."
+_DIFY_UI_ENVELOPE_KEY = "__dify_ui__"
+_MAX_UI_MESSAGE_PAYLOAD_BYTES = 128 * 1024
+_MAX_UI_MESSAGES_PER_TOOL_CALL = 16
+_MAX_UI_MESSAGES_PAYLOAD_BYTES = 512 * 1024
 
 
 class DifyCoreToolsDeps(LayerDeps):
@@ -104,13 +110,19 @@ class DifyCoreToolsLayer(PlainLayer[DifyCoreToolsDeps, DifyCoreToolsLayerConfig]
         tool_description = tool_config.description or tool_name
         tool_schema = deepcopy(tool_config.parameters_json_schema)
 
-        async def invoke_tool(_ctx: RunContext[object], **tool_arguments: object) -> str:
+        async def invoke_tool(_ctx: RunContext[object], **tool_arguments: object) -> str | ToolReturn:
             try:
                 response = await client.invoke(
                     execution_context=execution_context,
                     tool_config=tool_config,
                     tool_parameters=tool_arguments,
                 )
+                ui_messages = _extract_ui_messages(response.messages)
+                if ui_messages:
+                    return ToolReturn(
+                        response.observation,
+                        metadata={"dify_ui_messages": ui_messages},
+                    )
                 return response.observation
             except DifyCoreToolsClientConfigurationError:
                 return "Tool is unavailable because required execution context is missing."
@@ -141,6 +153,86 @@ class DifyCoreToolsLayer(PlainLayer[DifyCoreToolsDeps, DifyCoreToolsLayerConfig]
         )
 
 
+def _extract_ui_messages(messages: Sequence[Mapping[str, object]]) -> list[dict[str, object]]:
+    """Extract a bounded UI batch, dropping the complete batch on invalid input."""
+    ui_messages: list[dict[str, object]] = []
+    ui_payload_size = 2  # Compact JSON encoding of an empty list.
+
+    def append_ui_message(value: object) -> bool:
+        nonlocal ui_payload_size
+        if len(ui_messages) >= _MAX_UI_MESSAGES_PER_TOOL_CALL:
+            return False
+        payload = _normalize_ui_payload(value)
+        if payload is None:
+            return False
+        try:
+            encoded = json.dumps(
+                payload,
+                ensure_ascii=False,
+                separators=(",", ":"),
+                allow_nan=False,
+            ).encode()
+        except (TypeError, ValueError):
+            return False
+        if len(encoded) > _MAX_UI_MESSAGE_PAYLOAD_BYTES:
+            return False
+        candidate_size = ui_payload_size + len(encoded) + int(bool(ui_messages))
+        if candidate_size > _MAX_UI_MESSAGES_PAYLOAD_BYTES:
+            return False
+        ui_messages.append(payload)
+        ui_payload_size = candidate_size
+        return True
+
+    for tool_message in messages:
+        message_type = tool_message.get("type")
+        payload = tool_message.get("message")
+        if not isinstance(payload, dict):
+            continue
+        if message_type == "ui":
+            if not append_ui_message(payload):
+                return []
+            continue
+        if message_type == "variable":
+            if payload.get("variable_name") == _DIFY_UI_ENVELOPE_KEY and not append_ui_message(
+                payload.get("variable_value")
+            ):
+                return []
+            continue
+        if message_type == "json":
+            json_object = payload.get("json_object")
+            if isinstance(json_object, dict) and set(json_object) == {_DIFY_UI_ENVELOPE_KEY}:
+                if not append_ui_message(json_object[_DIFY_UI_ENVELOPE_KEY]):
+                    return []
+    return ui_messages
+
+
+def _normalize_ui_payload(value: object) -> dict[str, object] | None:
+    """Validate the transport-level UI envelope shared with the plugin boundary."""
+    if not isinstance(value, dict):
+        return None
+    required_keys = {"protocol", "protocol_version", "messages"}
+    allowed_keys = required_keys | {"fallback"}
+    if not required_keys.issubset(value) or not set(value).issubset(allowed_keys):
+        return None
+    if value.get("protocol") != "a2ui" or value.get("protocol_version") != "v0.9.1":
+        return None
+    ui_operations = value.get("messages")
+    if (
+        not isinstance(ui_operations, list)
+        or not 1 <= len(ui_operations) <= 64
+        or not all(isinstance(operation, dict) for operation in ui_operations)
+    ):
+        return None
+    fallback = value.get("fallback")
+    if fallback is not None and (not isinstance(fallback, str) or len(fallback) > 4096):
+        return None
+
+    payload = dict(value)
+    if fallback is None:
+        payload.pop("fallback", None)
+    return payload
+
+
 def _tool_error_text(*, tool_name: str, error: DifyCoreToolsClientError) -> str:
     if error.retryable:
         return TEMPORARY_UNAVAILABLE_OBSERVATION

dify-agent/src/dify_agent/layers/dify_plugin/tool_client.py

@@ -17,12 +17,14 @@ each tool's own ``plugin_id`` determines the transport identity placed in
 from __future__ import annotations
 
 import base64
+import json
 from collections.abc import AsyncIterator, Mapping
 from dataclasses import dataclass, field
 from enum import StrEnum
+from typing import Literal
 
 import httpx
-from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator, model_validator
 
 from dify_agent.layers.dify_plugin.configs import DifyPluginToolCredentialType
 from dify_agent.plugin_daemon_transport import (
@@ -62,6 +64,32 @@ class DifyPluginToolInvokeMessage(BaseModel):
         json_object: dict[str, object] | list[object]
         suppress_output: bool = False
 
+    class UIMessage(BaseModel):
+        protocol: Literal["a2ui"]
+        protocol_version: Literal["v0.9.1"]
+        messages: list[dict[str, object]] = Field(min_length=1, max_length=64)
+        fallback: str | None = Field(default=None, max_length=4096)
+
+        model_config = ConfigDict(extra="forbid", allow_inf_nan=False)
+
+        @model_validator(mode="after")
+        def validate_payload_size(self) -> "DifyPluginToolInvokeMessage.UIMessage":
+            payload = self.model_dump(mode="python", exclude_none=True)
+            try:
+                encoded = json.dumps(
+                    payload,
+                    ensure_ascii=False,
+                    separators=(",", ":"),
+                    allow_nan=False,
+                ).encode()
+            except TypeError as exc:
+                raise ValueError("UI message payload must be JSON-serializable") from exc
+            except ValueError as exc:
+                raise ValueError("UI message payload contains non-finite numbers") from exc
+            if len(encoded) > 131_072:
+                raise ValueError("UI message payload exceeds 128 KiB")
+            return self
+
     class BlobMessage(BaseModel):
         blob: bytes
 
@@ -108,10 +136,19 @@ class DifyPluginToolInvokeMessage(BaseModel):
         FILE = "file"
         LOG = "log"
         BLOB_CHUNK = "blob_chunk"
+        UI = "ui"
 
     type: MessageType = MessageType.TEXT
     message: (
-        TextMessage | JsonMessage | BlobChunkMessage | BlobMessage | LogMessage | FileMessage | VariableMessage | None
+        TextMessage
+        | JsonMessage
+        | UIMessage
+        | BlobChunkMessage
+        | BlobMessage
+        | LogMessage
+        | FileMessage
+        | VariableMessage
+        | None
     )
     meta: dict[str, object] | None = None
 

dify-agent/src/dify_agent/layers/dify_plugin/tools_layer.py

@@ -23,7 +23,7 @@ from urllib.parse import urlparse
 
 import httpx
 from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError
-from pydantic_ai import RunContext, Tool
+from pydantic_ai import RunContext, Tool, ToolReturn
 from pydantic_ai.tools import ToolDefinition
 from typing_extensions import Self, override
 
@@ -56,6 +56,9 @@ _FILE_UPLOAD_BEGIN = "<<<DIFY_PLUGIN_TOOL_FILE_UPLOAD_BEGIN>>>"
 _FILE_UPLOAD_END = "<<<DIFY_PLUGIN_TOOL_FILE_UPLOAD_END>>>"
 _FILE_UPLOAD_TIMEOUT_SECONDS = 60.0
 _SUPPORTED_REMOTE_URL_PREFIXES = ("http://", "https://")
+_DIFY_UI_ENVELOPE_KEY = "__dify_ui__"
+_MAX_UI_MESSAGES_PER_TOOL_CALL = 16
+_MAX_UI_MESSAGES_PAYLOAD_BYTES = 512 * 1024
 
 
 class DifyPluginToolsDeps(LayerDeps):
@@ -259,7 +262,7 @@ def _build_pydantic_ai_tool(
     tool_description = tool_config.description or tool_name
     tool_schema = deepcopy(tool_config.parameters_json_schema)
 
-    async def invoke_tool(_ctx: RunContext[object], **tool_arguments: object) -> str:
+    async def invoke_tool(_ctx: RunContext[object], **tool_arguments: object) -> str | ToolReturn:
         try:
             merged_arguments = await _prepare_tool_arguments(
                 effective_parameters,
@@ -274,7 +277,13 @@ def _build_pydantic_ai_tool(
                 credentials=dict(tool_config.credentials),
                 tool_parameters=merged_arguments,
             )
-            return _convert_tool_response_to_text(messages)
+            observation, ui_messages = _convert_tool_response(messages)
+            if ui_messages:
+                return ToolReturn(
+                    observation,
+                    metadata={"dify_ui_messages": ui_messages},
+                )
+            return observation
         except DifyPluginToolClientError as exc:
             return _tool_error_text(tool_name=tool_name, error=exc)
         except ValueError as exc:
@@ -643,8 +652,56 @@ def _convert_tool_response_to_text(tool_response: Sequence[DifyPluginToolInvokeM
     are deduplicated against existing text so mixed text/JSON streams do not
     repeat the same content unnecessarily.
     """
+    observation, _ = _convert_tool_response(tool_response)
+    return observation
+
+
+def _convert_tool_response(
+    tool_response: Sequence[DifyPluginToolInvokeMessage],
+) -> tuple[str, list[dict[str, object]]]:
+    """Separate observations from a bounded, fail-closed batch of UI messages.
+
+    Reserved compatibility envelopes never become model-visible JSON. A
+    malformed envelope, a seventeenth UI message, or a batch larger than 512
+    KiB invalidates the complete UI batch while text observation processing
+    continues.
+    """
     parts: list[str] = []
     json_parts: list[str] = []
+    ui_messages: list[dict[str, object]] = []
+    ui_payload_size = 2  # Compact JSON encoding of an empty list.
+    invalid_ui_batch = False
+
+    def append_ui_message(value: object) -> None:
+        nonlocal invalid_ui_batch, ui_payload_size
+        if invalid_ui_batch:
+            return
+        if len(ui_messages) >= _MAX_UI_MESSAGES_PER_TOOL_CALL:
+            invalid_ui_batch = True
+            ui_messages.clear()
+            return
+
+        try:
+            ui_message = DifyPluginToolInvokeMessage.UIMessage.model_validate(value)
+            payload = ui_message.model_dump(mode="json", exclude_none=True)
+            encoded = json.dumps(
+                payload,
+                ensure_ascii=False,
+                separators=(",", ":"),
+                allow_nan=False,
+            ).encode()
+        except (TypeError, ValueError):
+            invalid_ui_batch = True
+            ui_messages.clear()
+            return
+
+        candidate_size = ui_payload_size + len(encoded) + int(bool(ui_messages))
+        if candidate_size > _MAX_UI_MESSAGES_PAYLOAD_BYTES:
+            invalid_ui_batch = True
+            ui_messages.clear()
+            return
+        ui_messages.append(payload)
+        ui_payload_size = candidate_size
 
     for response in tool_response:
         if response.type is DifyPluginToolInvokeMessage.MessageType.TEXT:
@@ -665,9 +722,22 @@ def _convert_tool_response_to_text(tool_response: Sequence[DifyPluginToolInvokeM
             )
         elif response.type is DifyPluginToolInvokeMessage.MessageType.JSON:
             json_message = response.message
-            if isinstance(json_message, DifyPluginToolInvokeMessage.JsonMessage) and not json_message.suppress_output:
-                json_parts.append(json.dumps(json_message.json_object, ensure_ascii=False, default=str))
+            if isinstance(json_message, DifyPluginToolInvokeMessage.JsonMessage):
+                json_object = json_message.json_object
+                if isinstance(json_object, dict) and set(json_object) == {_DIFY_UI_ENVELOPE_KEY}:
+                    append_ui_message(json_object[_DIFY_UI_ENVELOPE_KEY])
+                    continue
+                if not json_message.suppress_output:
+                    json_parts.append(json.dumps(json_object, ensure_ascii=False, default=str))
+        elif response.type is DifyPluginToolInvokeMessage.MessageType.UI:
+            append_ui_message(response.message)
         elif response.type is DifyPluginToolInvokeMessage.MessageType.VARIABLE:
+            variable_message = response.message
+            if (
+                isinstance(variable_message, DifyPluginToolInvokeMessage.VariableMessage)
+                and variable_message.variable_name == _DIFY_UI_ENVELOPE_KEY
+            ):
+                append_ui_message(variable_message.variable_value)
             continue
         else:
             parts.append(str(response.message))
@@ -675,7 +745,7 @@ def _convert_tool_response_to_text(tool_response: Sequence[DifyPluginToolInvokeM
     if json_parts:
         existing_parts = set(parts)
         parts.extend(part for part in json_parts if part not in existing_parts)
-    return "".join(parts)
+    return "".join(parts), [] if invalid_ui_batch else ui_messages
 
 
 __all__ = ["DifyPluginToolsDeps", "DifyPluginToolsLayer"]

dify-agent/tests/local/dify_agent/layers/dify_core_tools/test_layer.py

@@ -1,9 +1,11 @@
 import asyncio
+import json
 import sys
 import types
 
 import httpx
 import pytest
+from pydantic_ai import ToolReturn
 
 from agenton.compositor import Compositor, LayerNode, LayerProvider
 
@@ -52,7 +54,7 @@ def _install_graphon_stubs() -> None:
 _install_graphon_stubs()
 
 from dify_agent.layers.dify_core_tools.configs import DifyCoreToolConfig, DifyCoreToolsLayerConfig  # noqa: E402
-from dify_agent.layers.dify_core_tools.layer import DifyCoreToolsLayer  # noqa: E402
+from dify_agent.layers.dify_core_tools.layer import DifyCoreToolsLayer, _extract_ui_messages  # noqa: E402
 from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig  # noqa: E402
 from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer  # noqa: E402
 
@@ -97,6 +99,54 @@ def _execution_context_config() -> DifyExecutionContextLayerConfig:
     )
 
 
+def _ui_payload(surface_id: str, *, text: object = "Sunny") -> dict[str, object]:
+    return {
+        "protocol": "a2ui",
+        "protocol_version": "v0.9.1",
+        "messages": [
+            {
+                "version": "v0.9.1",
+                "createSurface": {
+                    "surfaceId": surface_id,
+                    "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                },
+            },
+            {
+                "version": "v0.9.1",
+                "updateComponents": {
+                    "surfaceId": surface_id,
+                    "components": [{"id": "root", "component": "Text", "text": text}],
+                },
+            },
+        ],
+    }
+
+
+def _large_ui_payload(surface_id: str) -> dict[str, object]:
+    child_ids = [f"text-{index}" for index in range(23)]
+    payload = _ui_payload(surface_id)
+    messages = payload["messages"]
+    assert isinstance(messages, list)
+    update_components = messages[-1]["updateComponents"]  # type: ignore[index]
+    assert isinstance(update_components, dict)
+    update_components["components"] = [
+        {"id": "root", "component": "Column", "children": child_ids},
+        *[
+            {
+                "id": child_id,
+                "component": "Text",
+                "text": "x" * 4096,
+            }
+            for child_id in child_ids
+        ],
+    ]
+    return payload
+
+
+def _ui_tool_message(payload: dict[str, object]) -> dict[str, object]:
+    return {"type": "ui", "message": payload}
+
+
 def test_core_tools_layer_exposes_pydantic_ai_tool_and_returns_inner_api_observation() -> None:
     async def scenario() -> None:
         compositor = Compositor(
@@ -143,6 +193,117 @@ def test_core_tools_layer_exposes_pydantic_ai_tool_and_returns_inner_api_observa
     asyncio.run(scenario())
 
 
+def test_core_tools_layer_returns_ui_as_tool_metadata() -> None:
+    ui_message = _ui_payload("weather")
+
+    async def scenario() -> None:
+        compositor = Compositor(
+            [
+                LayerNode("execution_context", _execution_context_provider()),
+                LayerNode("core_tools", _core_tools_provider(), deps={"execution_context": "execution_context"}),
+            ]
+        )
+        async with httpx.AsyncClient(
+            transport=httpx.MockTransport(
+                lambda _request: httpx.Response(
+                    200,
+                    json={
+                        "messages": [{"type": "ui", "message": ui_message}],
+                        "observation": "Sunny",
+                        "metadata": {"provider_type": "builtin"},
+                    },
+                )
+            )
+        ) as http_client:
+            async with compositor.enter(
+                configs={
+                    "execution_context": _execution_context_config(),
+                    "core_tools": DifyCoreToolsLayerConfig(
+                        tools=[
+                            DifyCoreToolConfig(
+                                provider_type="builtin",
+                                provider_id="weather",
+                                tool_name="forecast",
+                                parameters=[],
+                                parameters_json_schema={"type": "object", "properties": {}, "required": []},
+                            )
+                        ]
+                    ),
+                }
+            ) as run:
+                layer = run.get_layer("core_tools", DifyCoreToolsLayer)
+                tool = (await layer.get_tools(http_client=http_client))[0]
+                result = await tool.function_schema.call({}, None)  # pyright: ignore[reportArgumentType]
+
+        assert isinstance(result, ToolReturn)
+        assert result.return_value == "Sunny"
+        assert result.metadata == {"dify_ui_messages": [ui_message]}
+
+    asyncio.run(scenario())
+
+
+def test_extract_ui_messages_allows_sixteen_and_drops_seventeenth() -> None:
+    payloads = [_ui_payload(f"surface-{index}") for index in range(17)]
+
+    assert _extract_ui_messages([_ui_tool_message(payload) for payload in payloads[:16]]) == payloads[:16]
+    assert _extract_ui_messages([_ui_tool_message(payload) for payload in payloads]) == []
+
+
+def test_extract_ui_messages_drops_aggregate_payload_overflow() -> None:
+    payloads = [_large_ui_payload(f"surface-{index}") for index in range(6)]
+    assert len(json.dumps(payloads, ensure_ascii=False, separators=(",", ":")).encode()) > 512 * 1024
+
+    assert _extract_ui_messages([_ui_tool_message(payload) for payload in payloads]) == []
+
+
+def test_extract_ui_messages_drops_single_payload_over_128_kib() -> None:
+    payload = _ui_payload("oversized", text="x" * 131_000)
+
+    assert _extract_ui_messages([_ui_tool_message(payload)]) == []
+
+
+@pytest.mark.parametrize(
+    "reserved_message",
+    [
+        {
+            "type": "json",
+            "message": {"json_object": {"__dify_ui__": "invalid"}},
+        },
+        {
+            "type": "json",
+            "message": {"json_object": {"__dify_ui__": {"protocol": "invalid"}}},
+        },
+        {
+            "type": "variable",
+            "message": {"variable_name": "__dify_ui__", "variable_value": "invalid"},
+        },
+    ],
+    ids=["json-non-object", "json-invalid-object", "variable-non-object"],
+)
+def test_extract_ui_messages_drops_batch_with_invalid_reserved_envelope(
+    reserved_message: dict[str, object],
+) -> None:
+    assert _extract_ui_messages([_ui_tool_message(_ui_payload("valid")), reserved_message]) == []
+
+
+@pytest.mark.parametrize(
+    "value",
+    [float("nan"), float("inf"), float("-inf")],
+    ids=["nan", "positive-infinity", "negative-infinity"],
+)
+def test_extract_ui_messages_rejects_non_finite_ui_without_restricting_ordinary_json(
+    value: float,
+) -> None:
+    assert _extract_ui_messages([_ui_tool_message(_ui_payload("invalid", text=value))]) == []
+
+    valid_payload = _ui_payload("valid")
+    ordinary_json = {
+        "type": "json",
+        "message": {"json_object": {"value": value}},
+    }
+    assert _extract_ui_messages([ordinary_json, _ui_tool_message(valid_payload)]) == [valid_payload]
+
+
 @pytest.mark.parametrize(
     ("response", "expected"),
     [

dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py

@@ -1,10 +1,12 @@
 import asyncio
 import json
+import math
 from types import SimpleNamespace
 
 import httpx
 import pytest
-from pydantic import JsonValue
+from pydantic import JsonValue, ValidationError
+from pydantic_ai import ToolReturn
 
 from agenton.compositor import Compositor, LayerNode, LayerProvider
 from dify_agent.adapters.llm import DifyLLMAdapterModel
@@ -20,7 +22,12 @@ from dify_agent.layers.dify_plugin.configs import (
     DifyPluginToolsLayerConfig,
 )
 from dify_agent.layers.dify_plugin.llm_layer import DifyPluginLLMLayer
-from dify_agent.layers.dify_plugin.tools_layer import DifyPluginToolsLayer, _PluginToolFileContext
+from dify_agent.layers.dify_plugin.tool_client import DifyPluginToolInvokeMessage
+from dify_agent.layers.dify_plugin.tools_layer import (
+    DifyPluginToolsLayer,
+    _convert_tool_response,
+    _PluginToolFileContext,
+)
 from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
 from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer
 
@@ -203,6 +210,7 @@ def _invoke_stream_response(
     *,
     error_payload: dict[str, object] | None = None,
     chunked_blob: bool = False,
+    ui_message: bool = False,
 ) -> httpx.Response:
     if error_payload is not None:
         return httpx.Response(400, json=error_payload)
@@ -217,6 +225,43 @@ def _invoke_stream_response(
         )
         return httpx.Response(200, text=stream_payload)
 
+    if ui_message:
+        ui_payload = {
+            "protocol": "a2ui",
+            "protocol_version": "v0.9.1",
+            "messages": [
+                {
+                    "version": "v0.9.1",
+                    "createSurface": {
+                        "surfaceId": "weather",
+                        "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                    },
+                },
+                {
+                    "version": "v0.9.1",
+                    "updateComponents": {
+                        "surfaceId": "weather",
+                        "components": [
+                            {
+                                "id": "root",
+                                "component": "Text",
+                                "text": "Sunny",
+                            }
+                        ],
+                    },
+                },
+            ],
+            "fallback": "Sunny",
+        }
+        stream_payload = "\n".join(
+            [
+                f"data: {json.dumps({'code': 0, 'message': 'ok', 'data': {'type': 'text', 'message': {'text': 'Sunny'}}})}",
+                f"data: {json.dumps({'code': 0, 'message': 'ok', 'data': {'type': 'ui', 'message': ui_payload}})}",
+                "",
+            ]
+        )
+        return httpx.Response(200, text=stream_payload)
+
     stream_payload = "\n".join(
         [
             f"data: {json.dumps({'code': 0, 'message': 'ok', 'data': {'type': 'text', 'message': {'text': 'found '}}})}",
@@ -231,6 +276,7 @@ def _tool_transport(
     *,
     invoke_error_payload: dict[str, object] | None = None,
     chunked_blob: bool = False,
+    ui_message: bool = False,
 ) -> httpx.MockTransport:
     def handler(request: httpx.Request) -> httpx.Response:
         if request.url.path.endswith("/dispatch/tool/invoke"):
@@ -245,13 +291,68 @@ def _tool_transport(
                 "api_version": "2026-01",
                 "auth_scope": "workspace",
             }
-            return _invoke_stream_response(error_payload=invoke_error_payload, chunked_blob=chunked_blob)
+            return _invoke_stream_response(
+                error_payload=invoke_error_payload,
+                chunked_blob=chunked_blob,
+                ui_message=ui_message,
+            )
 
         raise AssertionError(f"Unexpected request path: {request.url.path}")
 
     return httpx.MockTransport(handler)
 
 
+def _ui_payload(surface_id: str, *, text: object = "Sunny") -> dict[str, object]:
+    return {
+        "protocol": "a2ui",
+        "protocol_version": "v0.9.1",
+        "messages": [
+            {
+                "version": "v0.9.1",
+                "createSurface": {
+                    "surfaceId": surface_id,
+                    "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                },
+            },
+            {
+                "version": "v0.9.1",
+                "updateComponents": {
+                    "surfaceId": surface_id,
+                    "components": [{"id": "root", "component": "Text", "text": text}],
+                },
+            },
+        ],
+    }
+
+
+def _large_ui_payload(surface_id: str) -> dict[str, object]:
+    child_ids = [f"text-{index}" for index in range(23)]
+    payload = _ui_payload(surface_id)
+    messages = payload["messages"]
+    assert isinstance(messages, list)
+    update_components = messages[-1]["updateComponents"]  # type: ignore[index]
+    assert isinstance(update_components, dict)
+    update_components["components"] = [
+        {"id": "root", "component": "Column", "children": child_ids},
+        *[
+            {
+                "id": child_id,
+                "component": "Text",
+                "text": "x" * 4096,
+            }
+            for child_id in child_ids
+        ],
+    ]
+    return payload
+
+
+def _native_ui_message(payload: dict[str, object]) -> DifyPluginToolInvokeMessage:
+    return DifyPluginToolInvokeMessage(
+        type=DifyPluginToolInvokeMessage.MessageType.UI,
+        message=DifyPluginToolInvokeMessage.UIMessage.model_validate(payload),
+    )
+
+
 def _file_tool_transport(
     *,
     expected_source: object,
@@ -341,6 +442,188 @@ def test_dify_plugin_tools_layer_uses_prepared_tool_definition_and_invokes_daemo
     asyncio.run(scenario())
 
 
+def test_dify_plugin_tools_layer_keeps_ui_out_of_model_observation() -> None:
+    async def scenario() -> None:
+        compositor = Compositor(
+            [
+                LayerNode("execution_context", _execution_context_provider()),
+                LayerNode("tools", _tools_provider(), deps={"execution_context": "execution_context"}),
+            ]
+        )
+        async with httpx.AsyncClient(transport=_tool_transport(ui_message=True)) as client:
+            async with compositor.enter(
+                configs={"execution_context": _execution_context_config(), "tools": _tools_config()}
+            ) as run:
+                tool = (
+                    await run.get_layer("tools", DifyPluginToolsLayer).get_tools(
+                        http_client=client,
+                        dify_api_http_client=client,
+                    )
+                )[0]
+                result = await tool.function_schema.call(
+                    {"query": "dify", "region": "global"},
+                    None,  # pyright: ignore[reportArgumentType]
+                )
+
+        assert isinstance(result, ToolReturn)
+        assert result.return_value == "Sunny"
+        assert result.metadata == {
+            "dify_ui_messages": [
+                {
+                    "protocol": "a2ui",
+                    "protocol_version": "v0.9.1",
+                    "messages": [
+                        {
+                            "version": "v0.9.1",
+                            "createSurface": {
+                                "surfaceId": "weather",
+                                "catalogId": "https://dify.ai/a2ui/catalog/v1",
+                            },
+                        },
+                        {
+                            "version": "v0.9.1",
+                            "updateComponents": {
+                                "surfaceId": "weather",
+                                "components": [
+                                    {
+                                        "id": "root",
+                                        "component": "Text",
+                                        "text": "Sunny",
+                                    }
+                                ],
+                            },
+                        },
+                    ],
+                    "fallback": "Sunny",
+                }
+            ]
+        }
+
+    asyncio.run(scenario())
+
+
+def test_dify_plugin_tools_layer_supports_safe_variable_ui_fallback() -> None:
+    ui_message = _ui_payload("weather")
+    observation, ui_messages = _convert_tool_response(
+        [
+            DifyPluginToolInvokeMessage(
+                type=DifyPluginToolInvokeMessage.MessageType.VARIABLE,
+                message=DifyPluginToolInvokeMessage.VariableMessage(
+                    variable_name="__dify_ui__",
+                    variable_value=ui_message,
+                ),
+            )
+        ]
+    )
+
+    assert observation == ""
+    assert ui_messages == [ui_message]
+
+
+def test_dify_plugin_ui_message_payload_size_includes_full_envelope() -> None:
+    with pytest.raises(ValidationError, match="128 KiB"):
+        DifyPluginToolInvokeMessage.UIMessage(
+            protocol="a2ui",
+            protocol_version="v0.9.1",
+            messages=[{"value": "x" * 127_000}],
+            fallback="y" * 4096,
+        )
+
+
+def test_dify_plugin_ui_message_rejects_non_json_values() -> None:
+    with pytest.raises(ValidationError, match="JSON-serializable"):
+        DifyPluginToolInvokeMessage.UIMessage(
+            protocol="a2ui",
+            protocol_version="v0.9.1",
+            messages=[{"value": object()}],
+        )
+
+
+@pytest.mark.parametrize(
+    "value",
+    [float("nan"), float("inf"), float("-inf")],
+    ids=["nan", "positive-infinity", "negative-infinity"],
+)
+def test_dify_plugin_ui_message_rejects_non_finite_numbers_without_restricting_json_messages(
+    value: float,
+) -> None:
+    with pytest.raises(ValidationError, match="finite"):
+        DifyPluginToolInvokeMessage.UIMessage(
+            protocol="a2ui",
+            protocol_version="v0.9.1",
+            messages=[{"value": value}],
+        )
+
+    json_message = DifyPluginToolInvokeMessage.JsonMessage(json_object={"value": value})
+    assert isinstance(json_message.json_object, dict)
+    stored_value = json_message.json_object["value"]
+    assert isinstance(stored_value, float)
+    assert not math.isfinite(stored_value)
+
+
+def test_dify_plugin_tools_layer_drops_ui_batch_on_seventeenth_message() -> None:
+    observation, ui_messages = _convert_tool_response(
+        [
+            DifyPluginToolInvokeMessage(
+                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
+                message=DifyPluginToolInvokeMessage.TextMessage(text="before"),
+            ),
+            *[_native_ui_message(_ui_payload(f"surface-{index}")) for index in range(17)],
+            DifyPluginToolInvokeMessage(
+                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
+                message=DifyPluginToolInvokeMessage.TextMessage(text=" after"),
+            ),
+        ]
+    )
+
+    assert observation == "before after"
+    assert ui_messages == []
+
+
+def test_dify_plugin_tools_layer_drops_ui_batch_on_aggregate_payload_overflow() -> None:
+    payloads = [_large_ui_payload(f"surface-{index}") for index in range(6)]
+    assert len(json.dumps(payloads, ensure_ascii=False, separators=(",", ":")).encode()) > 512 * 1024
+
+    observation, ui_messages = _convert_tool_response(
+        [
+            DifyPluginToolInvokeMessage(
+                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
+                message=DifyPluginToolInvokeMessage.TextMessage(text="weather observation"),
+            ),
+            *[_native_ui_message(payload) for payload in payloads],
+        ]
+    )
+
+    assert observation == "weather observation"
+    assert ui_messages == []
+
+
+@pytest.mark.parametrize(
+    "invalid_payload",
+    ["invalid", {"protocol": "invalid"}],
+    ids=["non-object", "invalid-object"],
+)
+def test_dify_plugin_tools_layer_suppresses_invalid_reserved_ui_envelopes(
+    invalid_payload: object,
+) -> None:
+    observation, ui_messages = _convert_tool_response(
+        [
+            DifyPluginToolInvokeMessage(
+                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
+                message=DifyPluginToolInvokeMessage.TextMessage(text="weather observation"),
+            ),
+            _native_ui_message(_ui_payload("valid")),
+            DifyPluginToolInvokeMessage(
+                type=DifyPluginToolInvokeMessage.MessageType.JSON,
+                message=DifyPluginToolInvokeMessage.JsonMessage(json_object={"__dify_ui__": invalid_payload}),
+            ),
+        ]
+    )
+
+    assert observation == "weather observation"
+    assert ui_messages == []
+
+
 def test_dify_plugin_tools_layer_uses_each_tool_plugin_id_for_transport() -> None:
     async def scenario() -> None:
         seen_requests: list[tuple[str, str, str, str]] = []

docker/envs/core-services/shared.env.example

@@ -16,21 +16,11 @@ CHECK_UPDATE_URL=https://updates.dify.ai
 OPENAI_API_BASE=https://api.openai.com/v1
 MIGRATION_ENABLED=true
 FILES_ACCESS_TIMEOUT=300
-# System Features
-MARKETPLACE_ENABLED=true
-ENABLE_EMAIL_CODE_LOGIN=false
-ENABLE_EMAIL_PASSWORD_LOGIN=true
-ENABLE_SOCIAL_OAUTH_LOGIN=false
 # Remove `collaboration` from COMPOSE_PROFILES to stop the dedicated websocket service.
 ENABLE_COLLABORATION_MODE=true
-ALLOW_REGISTER=false
-ALLOW_CREATE_WORKSPACE=false
-ENABLE_CHANGE_EMAIL=true
-ENABLE_TRIAL_APP=false
-ENABLE_EXPLORE_BANNER=false
+
+# Learn app feature toggle
 ENABLE_LEARN_APP=true
-ENABLE_STEP_BY_STEP_TOUR=false
-RBAC_ENABLED=false
 CELERY_BROKER_URL=redis://:difyai123456@redis:6379/1
 CELERY_TASK_ANNOTATIONS=null
 AZURE_BLOB_ACCOUNT_URL=https://<your_account_name>.blob.core.windows.net
@@ -100,8 +90,6 @@ WORKFLOW_LOG_CLEANUP_SPECIFIC_WORKFLOW_IDS=
 EXPOSE_PLUGIN_DEBUGGING_HOST=localhost
 EXPOSE_PLUGIN_DEBUGGING_PORT=5003
 DEPLOY_ENV=PRODUCTION
-EDITION=SELF_HOSTED
-ENTERPRISE_ENABLED=false
 ACCESS_TOKEN_EXPIRE_MINUTES=60
 REFRESH_TOKEN_EXPIRE_DAYS=30
 APP_DEFAULT_ACTIVE_REQUESTS=0

docker/envs/core-services/web.env.example

@@ -15,6 +15,7 @@ TEXT_GENERATION_TIMEOUT_MS=60000
 ALLOW_INLINE_STYLES=false
 ALLOW_UNSAFE_DATA_SCHEME=false
 MAX_TREE_DEPTH=50
+MARKETPLACE_ENABLED=true
 MARKETPLACE_API_URL=https://marketplace.dify.ai
 INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH=4000
 ALLOW_EMBED=false

packages/contracts/generated/api/console/oauth/types.gen.ts

@@ -40,6 +40,7 @@ export type OAuthClientPayload = {
 export type OAuthProviderAccountResponse = {
   avatar?: string | null
   email: string
+  id: string
   interface_language: string
   name: string
   timezone: string

packages/contracts/generated/api/console/oauth/zod.gen.ts

@@ -60,6 +60,7 @@ export const zOAuthClientPayload = z.object({
 export const zOAuthProviderAccountResponse = z.object({
   avatar: z.string().nullish(),
   email: z.string(),
+  id: z.string(),
   interface_language: z.string(),
   name: z.string(),
   timezone: z.string(),

packages/contracts/generated/api/console/system-features/types.gen.ts

@@ -6,7 +6,6 @@ export type ClientOptions = {
 
 export type SystemFeatureModel = {
   branding: BrandingModel
-  deployment_edition: DeploymentEdition
   enable_app_deploy: boolean
   enable_change_email: boolean
   enable_collaboration_mode: boolean
@@ -41,8 +40,6 @@ export type BrandingModel = {
   workspace_logo: string
 }
 
-export type DeploymentEdition = 'CLOUD' | 'COMMUNITY' | 'ENTERPRISE'
-
 export type LicenseModel = {
   expired_at: string
   seats: LicenseLimitationModel

packages/contracts/generated/api/console/system-features/zod.gen.ts

@@ -13,11 +13,6 @@ export const zBrandingModel = z.object({
   workspace_logo: z.string().default(''),
 })
 
-/**
- * DeploymentEdition
- */
-export const zDeploymentEdition = z.enum(['CLOUD', 'COMMUNITY', 'ENTERPRISE'])
-
 /**
  * PluginManagerModel
  */
@@ -109,7 +104,6 @@ export const zSystemFeatureModel = z.object({
     login_page_logo: '',
     workspace_logo: '',
   }),
-  deployment_edition: zDeploymentEdition,
   enable_app_deploy: z.boolean().default(false),
   enable_change_email: z.boolean().default(true),
   enable_collaboration_mode: z.boolean().default(true),

packages/contracts/generated/api/web/types.gen.ts

@@ -118,8 +118,6 @@ export type ConversationRenamePayload = (
   name?: string | null
 }
 
-export type DeploymentEdition = 'CLOUD' | 'COMMUNITY' | 'ENTERPRISE'
-
 export type EmailCodeLoginSendPayload = {
   email: string
   language?: string | null
@@ -512,7 +510,6 @@ export type SuggestedQuestionsResponse = {
 
 export type SystemFeatureModel = {
   branding: BrandingModel
-  deployment_edition: DeploymentEdition
   enable_app_deploy: boolean
   enable_change_email: boolean
   enable_collaboration_mode: boolean

packages/contracts/generated/api/web/zod.gen.ts

@@ -132,11 +132,6 @@ export const zConversationRenamePayload = z.intersection(
   }),
 )
 
-/**
- * DeploymentEdition
- */
-export const zDeploymentEdition = z.enum(['CLOUD', 'COMMUNITY', 'ENTERPRISE'])
-
 /**
  * EmailCodeLoginSendPayload
  */
@@ -787,7 +782,6 @@ export const zSystemFeatureModel = z.object({
     login_page_logo: '',
     workspace_logo: '',
   }),
-  deployment_edition: zDeploymentEdition,
   enable_app_deploy: z.boolean().default(false),
   enable_change_email: z.boolean().default(true),
   enable_collaboration_mode: z.boolean().default(true),

packages/iconify-collections/assets/public/common/confluence.svg

@@ -0,0 +1,4 @@
+<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
+  <path fill="#1868DB" d="M4.1 15.7c-.4.7-.8 1.5-1.1 2.1-.2.5 0 1 .5 1.3l3.4 1.6c.5.2 1 0 1.3-.4.3-.6.7-1.3 1.1-2 2.9-4.8 5.9-4.2 11.1-1.8.5.2 1.1 0 1.3-.5l1.3-3.5c.2-.5-.1-1.1-.6-1.3-7.6-3.5-13.8-3.6-18.3 4.5Z"/>
+  <path fill="#1868DB" fill-opacity=".55" d="M19.9 8.3c.4-.7.8-1.5 1.1-2.1.2-.5 0-1-.5-1.3l-3.4-1.6c-.5-.2-1 0-1.3.4-.3.6-.7 1.3-1.1 2-2.9 4.8-5.9 4.2-11.1 1.8-.5-.2-1.1 0-1.3.5L1 11.5c-.2.5.1 1.1.6 1.3 7.6 3.5 13.8 3.6 18.3-4.5Z"/>
+</svg>

packages/iconify-collections/custom-public/icons.json

@@ -1,6 +1,6 @@
 {
   "prefix": "custom-public",
-  "lastModified": 1783583299,
+  "lastModified": 1784810412,
   "icons": {
     "agent-building-blocks": {
       "body": "<path fill=\"#155AEF\" fill-rule=\"evenodd\" d=\"M8.303 1.546c.178-.045.364-.051.544-.017c.23.043.432.167.573.246l3.757 2.113c.12.067.29.156.433.289l.06.06c.12.131.21.288.267.457c.07.215.063.445.063.6V9.56c0 .146.007.36-.056.563q-.055.181-.162.338l-.075.1c-.137.163-.32.274-.442.353l-5.013 3.259c-.135.088-.33.224-.556.282a1.3 1.3 0 0 1-.543.017c-.23-.043-.433-.166-.573-.245l-3.757-2.114c-.136-.077-.34-.182-.493-.35a1.3 1.3 0 0 1-.267-.456C1.993 11.09 2 10.86 2 10.704V6.441c0-.146-.007-.36.055-.563l.043-.118a1.3 1.3 0 0 1 .195-.32l.053-.059c.128-.131.282-.225.389-.294L7.86 1.755c.122-.078.273-.165.443-.209m-4.97 9.158l.001.164l.033.02l.11.062l3.264 1.836v-1.137L3.333 9.732zm4.741.917v1.076l4.464-2.901l.098-.065l.029-.02v-.034l.001-.118v-.923zm-4.74-3.419L6.74 10.12V8.982L3.333 7.066zm4.74.752v1.076l4.592-2.985V5.969zm.51-6.08l-4.631 3.01l3.429 1.93l4.664-3.032l-3.28-1.846l-.15-.082z\" clip-rule=\"evenodd\"/>"
@@ -65,6 +65,11 @@
       "width": 13,
       "height": 13
     },
+    "common-confluence": {
+      "body": "<g fill=\"#1868DB\"><path d=\"M4.1 15.7c-.4.7-.8 1.5-1.1 2.1c-.2.5 0 1 .5 1.3l3.4 1.6c.5.2 1 0 1.3-.4c.3-.6.7-1.3 1.1-2c2.9-4.8 5.9-4.2 11.1-1.8c.5.2 1.1 0 1.3-.5l1.3-3.5c.2-.5-.1-1.1-.6-1.3c-7.6-3.5-13.8-3.6-18.3 4.5\"/><path fill-opacity=\".55\" d=\"M19.9 8.3c.4-.7.8-1.5 1.1-2.1c.2-.5 0-1-.5-1.3l-3.4-1.6c-.5-.2-1 0-1.3.4c-.3.6-.7 1.3-1.1 2c-2.9 4.8-5.9 4.2-11.1 1.8c-.5-.2-1.1 0-1.3.5L1 11.5c-.2.5.1 1.1.6 1.3c7.6 3.5 13.8 3.6 18.3-4.5\"/></g>",
+      "width": 24,
+      "height": 24
+    },
     "common-d": {
       "body": "<g fill=\"none\"><path fill=\"#fff\" d=\"M2 1h5.943a7 7 0 1 1 0 14H2z\"/><path fill=\"url(#svgID0)\" d=\"M2 1h5.943a7 7 0 1 1 0 14H2z\"/><path fill=\"url(#svgID1)\" d=\"M7.943 8h.265v7h-.265z\"/><defs><radialGradient id=\"svgID0\" cx=\"0\" cy=\"0\" r=\"1\" gradientTransform=\"matrix(0 8.75 -8.75 0 7.943 8)\" gradientUnits=\"userSpaceOnUse\"><stop stop-color=\"#001FC2\"/><stop offset=\".711\" stop-color=\"#0667F8\" stop-opacity=\".2\"/><stop offset=\"1\" stop-color=\"#155EEF\" stop-opacity=\"0\"/></radialGradient><linearGradient id=\"svgID1\" x1=\"8.062\" x2=\"7.937\" y1=\"8.438\" y2=\"9.203\" gradientUnits=\"userSpaceOnUse\"><stop stop-color=\"#fff\" stop-opacity=\"0\"/><stop offset=\"1\" stop-color=\"#fff\"/></linearGradient></defs></g>"
     },

packages/iconify-collections/custom-public/info.json

@@ -1,7 +1,7 @@
 {
   "prefix": "custom-public",
   "name": "Dify Custom Public",
-  "total": 146,
+  "total": 147,
   "version": "0.0.0-private",
   "author": {
     "name": "LangGenius, Inc.",

web/.env.example

@@ -1,5 +1,9 @@
 # For production release, change this to PRODUCTION
 NEXT_PUBLIC_DEPLOY_ENV=DEVELOPMENT
+# The deployment edition, SELF_HOSTED
+NEXT_PUBLIC_EDITION=SELF_HOSTED
+# Whether a self-hosted deployment runs Enterprise Edition
+NEXT_PUBLIC_ENTERPRISE_ENABLED=false
 # The base path for the application
 NEXT_PUBLIC_BASE_PATH=
 # Server-only console API origin for server-side requests.
@@ -110,3 +114,20 @@ NEXT_PUBLIC_WEB_PREFIX=
 
 # number of concurrency
 NEXT_PUBLIC_BATCH_CONCURRENCY=5
+
+# Cloud system-features frontend defaults.
+# These values are only used when NEXT_PUBLIC_EDITION=CLOUD (IS_CLOUD_EDITION).
+NEXT_PUBLIC_ENABLE_MARKETPLACE=true
+NEXT_PUBLIC_ENABLE_EMAIL_CODE_LOGIN=true
+NEXT_PUBLIC_ENABLE_EMAIL_PASSWORD_LOGIN=false
+NEXT_PUBLIC_ENABLE_SOCIAL_OAUTH_LOGIN=true
+NEXT_PUBLIC_ENABLE_COLLABORATION_MODE=false
+NEXT_PUBLIC_ALLOW_REGISTER=true
+NEXT_PUBLIC_ALLOW_CREATE_WORKSPACE=true
+NEXT_PUBLIC_IS_EMAIL_SETUP=true
+NEXT_PUBLIC_ENABLE_CHANGE_EMAIL=true
+NEXT_PUBLIC_CREATORS_PLATFORM_FEATURES_ENABLED=true
+NEXT_PUBLIC_ENABLE_TRIAL_APP=true
+NEXT_PUBLIC_ENABLE_EXPLORE_BANNER=true
+NEXT_PUBLIC_RBAC_ENABLED=false
+NEXT_PUBLIC_KNOWLEDGE_FS_ENABLED=false

web/Dockerfile

@@ -49,6 +49,7 @@ RUN pnpm build && pnpm build:vinext
 FROM base AS production
 
 ENV NODE_ENV=production
+ENV EDITION=SELF_HOSTED
 ENV DEPLOY_ENV=PRODUCTION
 ENV CONSOLE_API_URL=http://127.0.0.1:5001
 ENV APP_API_URL=http://127.0.0.1:5001

web/__tests__/billing/billing-integration.test.tsx

@@ -1,5 +1,3 @@
-import type { RenderOptions } from '@testing-library/react'
-import type { ReactElement } from 'react'
 import type { UsagePlanInfo, UsageResetInfo } from '@/app/components/billing/type'
 import { screen } from '@testing-library/react'
 import userEvent from '@testing-library/user-event'
@@ -16,21 +14,21 @@ import TriggerEventsLimitModal from '@/app/components/billing/trigger-events-lim
 import { Plan } from '@/app/components/billing/type'
 import UpgradeBtn from '@/app/components/billing/upgrade-btn'
 import VectorSpaceFull from '@/app/components/billing/vector-space-full'
-import { createConsoleQueryWrapper } from '@/test/console/query-data'
-import { render as renderWithConsoleState } from '@/test/console/render'
-
-const render = (ui: ReactElement, options: RenderOptions = {}) => {
-  const { wrapper } = createConsoleQueryWrapper({
-    systemFeatures: { deployment_edition: 'CLOUD' },
-  })
-  return renderWithConsoleState(ui, { ...options, wrapper })
-}
+import { render } from '@/test/console/render'
 
 let mockProviderCtx: Record<string, unknown> = {}
 let mockConsoleState: Record<string, unknown> = {}
 const mockSetShowPricingModal = vi.fn()
 const mockSetShowAccountSettingModal = vi.fn()
 
+vi.mock('@/config', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@/config')>()
+  return {
+    ...actual,
+    IS_CLOUD_EDITION: true,
+  }
+})
+
 vi.mock('@/context/provider-context', () => ({
   useProviderContext: () => mockProviderCtx,
 }))

web/__tests__/billing/education-verification-flow.test.tsx

@@ -1,5 +1,3 @@
-import type { RenderOptions } from '@testing-library/react'
-import type { ReactElement } from 'react'
 /**
  * Integration test: Education Verification Flow
  *
@@ -16,15 +14,7 @@ import * as React from 'react'
 import { defaultPlan } from '@/app/components/billing/config'
 import PlanComp from '@/app/components/billing/plan'
 import { Plan } from '@/app/components/billing/type'
-import { createConsoleQueryWrapper } from '@/test/console/query-data'
-import { render as renderWithConsoleState } from '@/test/console/render'
-
-const render = (ui: ReactElement, options: RenderOptions = {}) => {
-  const { wrapper } = createConsoleQueryWrapper({
-    systemFeatures: { deployment_edition: 'CLOUD' },
-  })
-  return renderWithConsoleState(ui, { ...options, wrapper })
-}
+import { render } from '@/test/console/render'
 
 // ─── Mock state ──────────────────────────────────────────────────────────────
 let mockProviderCtx: Record<string, unknown> = {}
@@ -35,6 +25,14 @@ const mockRouterPush = vi.fn()
 const mockMutateAsync = vi.fn()
 const mockSetEducationVerifying = vi.hoisted(() => vi.fn())
 
+vi.mock('@/config', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@/config')>()
+  return {
+    ...actual,
+    IS_CLOUD_EDITION: true,
+  }
+})
+
 // ─── Context mocks ───────────────────────────────────────────────────────────
 vi.mock('@/context/provider-context', () => ({
   useProviderContext: () => mockProviderCtx,

web/__tests__/billing/partner-stack-flow.test.tsx

@@ -7,15 +7,11 @@
  * Covers URL param reading, cookie persistence, API bind on mount,
  * cookie cleanup after successful bind, and error handling for 400 status.
  */
-import { act, cleanup, renderHook, waitFor } from '@testing-library/react'
+import { act, cleanup, render, renderHook, waitFor } from '@testing-library/react'
 import Cookies from 'js-cookie'
 import * as React from 'react'
 import usePSInfo from '@/app/components/billing/partner-stack/use-ps-info'
 import { PARTNER_STACK_CONFIG } from '@/config'
-import { renderWithConsoleQuery } from '@/test/console/query-data'
-
-const render = (ui: React.ReactElement) =>
-  renderWithConsoleQuery(ui, { systemFeatures: { deployment_edition: 'CLOUD' } })
 
 // ─── Mock state ──────────────────────────────────────────────────────────────
 let mockSearchParams = new URLSearchParams()
@@ -43,6 +39,7 @@ vi.mock('@/config', async (importOriginal) => {
   const actual = await importOriginal<Record<string, unknown>>()
   return {
     ...actual,
+    IS_CLOUD_EDITION: true,
     PARTNER_STACK_CONFIG: {
       cookieName: 'partner_stack_info',
       saveCookieDays: 90,
@@ -291,7 +288,7 @@ describe('Partner Stack Flow', () => {
 
   // ─── 4. PartnerStack Component Mount ────────────────────────────────────
   describe('PartnerStack component mount behavior', () => {
-    it('should call saveOrUpdate and bind on mount', async () => {
+    it('should call saveOrUpdate and bind on mount when IS_CLOUD_EDITION is true', async () => {
       mockSearchParams = new URLSearchParams({
         ps_partner_key: 'mount-partner',
         ps_xid: 'mount-click',
@@ -302,13 +299,18 @@ describe('Partner Stack Flow', () => {
 
       render(<PartnerStack />)
 
+      // The component calls saveOrUpdate and bind in useEffect
       await waitFor(() => {
+        // Bind should have been called
         expect(mockMutateAsync).toHaveBeenCalledWith({
           partnerKey: 'mount-partner',
           clickId: 'mount-click',
         })
-        expect(Cookies.get(PARTNER_STACK_CONFIG.cookieName)).toBeUndefined()
       })
+
+      // Cookie should have been saved (saveOrUpdate was called before bind)
+      // After bind succeeds, cookie is removed
+      expect(Cookies.get(PARTNER_STACK_CONFIG.cookieName)).toBeUndefined()
     })
 
     it('should render nothing (return null)', async () => {