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/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

@@ -19559,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 |

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/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/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/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/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/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:

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]] = []

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/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/app/(commonLayout)/datasets/layout.spec.tsx

@@ -202,39 +202,45 @@ describe('DatasetsLayout', () => {
     expect(mockReplace).not.toHaveBeenCalled()
   })
 
-  it('should redirect direct external dataset connection route to /datasets without dataset.external.connect', async () => {
-    mockPathname = '/datasets/connect'
-    setConsoleState({
-      workspacePermissionKeys: [],
-    })
+  it.each(['/datasets/connect', '/datasets/new/space-1/sources/new'])(
+    'should redirect direct external source route to /datasets without dataset.external.connect: %s',
+    async (pathname) => {
+      mockPathname = pathname
+      setConsoleState({
+        workspacePermissionKeys: [],
+      })
 
-    render(
-      <DatasetsLayout>
-        <div>datasets</div>
-      </DatasetsLayout>,
-    )
+      render(
+        <DatasetsLayout>
+          <div>datasets</div>
+        </DatasetsLayout>,
+      )
 
-    expect(screen.queryByText('datasets')).not.toBeInTheDocument()
-    await waitFor(() => {
-      expect(mockReplace).toHaveBeenCalledWith('/datasets')
-    })
-  })
+      expect(screen.queryByText('datasets')).not.toBeInTheDocument()
+      await waitFor(() => {
+        expect(mockReplace).toHaveBeenCalledWith('/datasets')
+      })
+    },
+  )
 
-  it('should render direct external dataset connection route when workspace has dataset.external.connect', () => {
-    mockPathname = '/datasets/connect'
-    setConsoleState({
-      workspacePermissionKeys: ['dataset.external.connect'],
-    })
+  it.each(['/datasets/connect', '/datasets/new/space-1/sources/new'])(
+    'should render direct external source route with dataset.external.connect: %s',
+    (pathname) => {
+      mockPathname = pathname
+      setConsoleState({
+        workspacePermissionKeys: ['dataset.external.connect'],
+      })
 
-    render(
-      <DatasetsLayout>
-        <div>datasets</div>
-      </DatasetsLayout>,
-    )
+      render(
+        <DatasetsLayout>
+          <div>datasets</div>
+        </DatasetsLayout>,
+      )
 
-    expect(screen.getByText('datasets')).toBeInTheDocument()
-    expect(mockReplace).not.toHaveBeenCalled()
-  })
+      expect(screen.getByText('datasets')).toBeInTheDocument()
+      expect(mockReplace).not.toHaveBeenCalled()
+    },
+  )
 
   it('should disable external knowledge API queries without dataset.external.connect', () => {
     setConsoleState({

web/app/(commonLayout)/datasets/layout.tsx

@@ -24,7 +24,11 @@ const isDatasetCreatePath = (pathname: string) => {
 }
 
 const isDatasetExternalConnectPath = (pathname: string) => {
-  return pathname === '/datasets/connect' || pathname.startsWith('/datasets/connect/')
+  return (
+    pathname === '/datasets/connect' ||
+    pathname.startsWith('/datasets/connect/') ||
+    /^\/datasets\/new\/[^/]+\/sources\/new\/?$/.test(pathname)
+  )
 }
 
 export default function DatasetsLayout({ children }: { children: React.ReactNode }) {

web/app/(commonLayout)/datasets/new/[knowledgeSpaceId]/sources/new/page.tsx

@@ -0,0 +1,7 @@
+import { AddSourcePage } from '@/features/new-rag/add-source-page'
+
+export default async function Page({ params }: { params: Promise<{ knowledgeSpaceId: string }> }) {
+  const { knowledgeSpaceId } = await params
+
+  return <AddSourcePage knowledgeSpaceId={knowledgeSpaceId} />
+}

web/app/(commonLayout)/datasets/new/[knowledgeSpaceId]/sources/page.tsx

@@ -1,5 +1,7 @@
-import { KnowledgeRoutePlaceholder } from '@/features/new-rag/knowledge-route-placeholder'
+import { SourcesPage } from '@/features/new-rag/sources-page'
 
-export default function Page() {
-  return <KnowledgeRoutePlaceholder type="sources" />
+export default async function Page({ params }: { params: Promise<{ knowledgeSpaceId: string }> }) {
+  const { knowledgeSpaceId } = await params
+
+  return <SourcesPage knowledgeSpaceId={knowledgeSpaceId} />
 }

web/app/account/oauth/authorize/__tests__/page.spec.tsx

@@ -63,7 +63,8 @@ describe('OAuthAuthorize', () => {
     vi.clearAllMocks()
     mocks.searchParams = new URLSearchParams({
       client_id: 'client-1',
-      redirect_uri: 'https://client.example.com/callback?state=state-1',
+      redirect_uri: 'https://client.example.com/callback',
+      state: 'state-1',
     })
     mocks.request.mockImplementation(async (url: string) => {
       if (url.endsWith('/oauth/provider/authorize')) return jsonResponse({ code: 'oauth-code' })
@@ -86,7 +87,7 @@ describe('OAuthAuthorize', () => {
     vi.unstubAllGlobals()
   })
 
-  it('authorizes the displayed app and redirects with the returned code', async () => {
+  it('authorizes the displayed app and redirects with the returned code and state', async () => {
     const user = userEvent.setup()
     renderPage()
 
@@ -95,7 +96,7 @@ describe('OAuthAuthorize', () => {
     const providerTransportRequest = providerRequest?.[2]?.request as Request
     await expect(providerTransportRequest.clone().json()).resolves.toEqual({
       client_id: 'client-1',
-      redirect_uri: 'https://client.example.com/callback?state=state-1',
+      redirect_uri: 'https://client.example.com/callback',
     })
 
     await user.click(screen.getByRole('button', { name: /continue/i }))
@@ -106,7 +107,7 @@ describe('OAuthAuthorize', () => {
     await expect(transportRequest.clone().json()).resolves.toEqual({ client_id: 'client-1' })
     await waitFor(() =>
       expect(globalThis.location.href).toBe(
-        'https://client.example.com/callback?state=state-1&code=oauth-code',
+        'https://client.example.com/callback?code=oauth-code&state=state-1',
       ),
     )
   })

web/app/account/oauth/authorize/page.tsx

@@ -64,6 +64,7 @@ export default function OAuthAuthorize() {
   const searchParams = useSearchParams()
   const client_id = decodeURIComponent(searchParams.get('client_id') || '')
   const redirect_uri = decodeURIComponent(searchParams.get('redirect_uri') || '')
+  const state = searchParams.get('state')
   const hasOAuthParams = Boolean(client_id && redirect_uri)
   // Probe user profile. 401 stays as `error` (legitimate "not logged in" state),
   // other errors throw to the nearest error.tsx; jumpTo same-pathname guard in
@@ -117,6 +118,7 @@ export default function OAuthAuthorize() {
       const { code } = await authorize({ body: { client_id } })
       const url = new URL(redirect_uri)
       url.searchParams.set('code', code)
+      if (state) url.searchParams.set('state', state)
       globalThis.location.href = url.toString()
     } catch (error: unknown) {
       const message = error instanceof Error ? error.message : String(error)

web/app/components/app/configuration/config/config-document.tsx

@@ -45,7 +45,7 @@ const ConfigDocument: FC = () => {
   if (!isShowDocumentConfig || (readonly && !isDocumentEnabled)) return null
 
   return (
-    <div className="mt-2 flex items-center gap-2 rounded-xl border-t-[0.5px] border-l-[0.5px] bg-background-section-burn p-2">
+    <div className="mt-2 flex items-center gap-2 rounded-xl border-t-[0.5px] border-l-[0.5px] border-effects-highlight bg-background-section-burn p-2">
       <div className="shrink-0 p-1">
         <div className="rounded-lg border-[0.5px] border-divider-subtle bg-util-colors-indigo-indigo-600 p-1 shadow-xs">
           <Document className="size-4 text-text-primary-on-surface" />
@@ -64,7 +64,7 @@ const ConfigDocument: FC = () => {
       </div>
       {!readonly && (
         <div className="flex shrink-0 items-center">
-          <div className="mr-3 ml-1 h-3.5 w-px bg-divider-subtle"></div>
+          <div className="mr-3 ml-1 h-3.5 w-px bg-divider-regular"></div>
           <Switch checked={isDocumentEnabled} onCheckedChange={handleChange} size="md" />
         </div>
       )}

web/app/components/app/configuration/debug/debug-with-single-model/__tests__/index.spec.tsx

@@ -16,6 +16,7 @@ import {
   ModelTypeEnum,
 } from '@/app/components/header/account-setting/model-provider-page/declarations'
 import { CollectionType } from '@/app/components/tools/types'
+import { SupportUploadFileTypes } from '@/app/components/workflow/types'
 import { PromptMode } from '@/models/debug'
 import { renderWithAccountProfile as render } from '@/test/console/account-profile'
 import { AgentStrategy, AppModeEnum, ModelModeType, Resolution, TransferMethod } from '@/types/app'
@@ -461,6 +462,16 @@ const mockFile: FileEntity = {
   supportFileType: 'image',
 }
 
+const mockDocumentFile: FileEntity = {
+  id: 'file-2',
+  name: 'test.pdf',
+  size: 456,
+  type: 'application/pdf',
+  progress: 100,
+  transferMethod: TransferMethod.local_file,
+  supportFileType: SupportUploadFileTypes.document,
+}
+
 // Mock Chat component (complex with many dependencies)
 // This is a pragmatic mock that tests the integration at DebugWithSingleModel level
 vi.mock('@/app/components/base/chat/chat', () => ({
@@ -519,6 +530,13 @@ vi.mock('@/app/components/base/chat/chat', () => ({
         >
           Send With Files
         </button>
+        <button
+          data-testid="send-with-document"
+          onClick={() => onSend?.('test message', [mockDocumentFile])}
+          disabled={isResponding || readonly || inputDisabled}
+        >
+          Send With Document
+        </button>
         {isResponding && (
           <button data-testid="stop-button" onClick={onStopResponding}>
             Stop
@@ -985,7 +1003,7 @@ describe('DebugWithSingleModel', () => {
 
   // File Upload Tests
   describe('File Upload', () => {
-    it('should not include files when vision is not supported', async () => {
+    it('should include document files when document is supported without vision', async () => {
       mockUseDebugConfigurationContext.mockReturnValue({
         ...mockDebugConfigContext,
         modelConfig: createMockModelConfig({
@@ -1006,7 +1024,7 @@ describe('DebugWithSingleModel', () => {
                   model: 'gpt-3.5-turbo',
                   label: { en_US: 'GPT-3.5', zh_Hans: 'GPT-3.5' },
                   model_type: ModelTypeEnum.textGeneration,
-                  features: [], // No vision
+                  features: [ModelFeatureEnum.document],
                   fetch_from: ConfigurationMethodEnum.predefinedModel,
                   model_properties: {},
                   deprecated: false,
@@ -1026,14 +1044,21 @@ describe('DebugWithSingleModel', () => {
 
       render(<DebugWithSingleModel ref={ref as RefObject<DebugWithSingleModelRefType>} />)
 
-      fireEvent.click(screen.getByTestId('send-with-files'))
+      fireEvent.click(screen.getByTestId('send-with-document'))
 
       await waitFor(() => {
         expect(mockSsePost).toHaveBeenCalled()
       })
 
       const body = mockSsePost.mock.calls[0]![1].body
-      expect(body.files).toEqual([])
+      expect(body.files).toEqual([
+        {
+          type: SupportUploadFileTypes.document,
+          transfer_method: TransferMethod.local_file,
+          url: '',
+          upload_file_id: '',
+        },
+      ])
     })
 
     it('should support files when vision is enabled', async () => {

web/app/components/app/configuration/debug/debug-with-single-model/index.tsx

@@ -9,7 +9,6 @@ import Chat from '@/app/components/base/chat/chat'
 import { useChat } from '@/app/components/base/chat/chat/hooks'
 import { getLastAnswer, isValidGeneratedAnswer } from '@/app/components/base/chat/utils'
 import { useFeatures } from '@/app/components/base/features/hooks'
-import { ModelFeatureEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
 import { userProfileAtom } from '@/context/account-state'
 import { useDebugConfigurationContext } from '@/context/debug-configuration'
 import { useProviderContext } from '@/context/provider-context'
@@ -97,13 +96,6 @@ const DebugWithSingleModel = ({
     (message, files, isRegenerate = false, parentAnswer: ChatItem | null = null) => {
       if (!canTestAndRun) return
       if (checkCanSend && !checkCanSend()) return
-      const currentProvider = textGenerationModelList.find(
-        (item) => item.provider === modelConfig.provider,
-      )
-      const currentModel = currentProvider?.models.find(
-        (model) => model.model === modelConfig.model_id,
-      )
-      const supportVision = currentModel?.features?.includes(ModelFeatureEnum.vision)
 
       const configData = {
         ...config,
@@ -122,7 +114,7 @@ const DebugWithSingleModel = ({
         parent_message_id: (isRegenerate ? parentAnswer?.id : getLastAnswer(chatList)?.id) || null,
       }
 
-      if ((config.file_upload as any)?.enabled && files?.length && supportVision) data.files = files
+      if ((config.file_upload as any)?.enabled && files?.length) data.files = files
 
       handleSend(`apps/${appId}/chat-messages`, data, {
         onGetConversationMessages: (conversationId, getAbortController) =>

web/app/components/base/chat/chat-with-history/__tests__/hooks.spec.tsx

@@ -2,6 +2,7 @@ import type { ReactNode } from 'react'
 import type { ChatConfig } from '../../types'
 import type { InstalledApp } from '@/models/explore'
 import type { AppConversationData, AppData, AppMeta, ConversationItem } from '@/models/share'
+import type { UIPart } from '@/types/a2ui'
 import { ToastHost } from '@langgenius/dify-ui/toast'
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import { act, renderHook, waitFor } from '@testing-library/react'
@@ -146,6 +147,29 @@ const setConversationIdInfo = (appId: string, conversationId: string) => {
   localStorage.setItem(CONVERSATION_ID_INFO, JSON.stringify(value))
 }
 
+const createHistoricalUIPart = (): UIPart => ({
+  part_id: 'call/abc=:weather',
+  sequence: 1,
+  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: 'Card', children: [], title: 'Weather' }],
+      },
+    },
+  ],
+})
+
 // Scenario: useChatWithHistory integrates share queries for conversations and chat list.
 describe('useChatWithHistory', () => {
   beforeEach(() => {
@@ -1094,6 +1118,32 @@ describe('useChatWithHistory', () => {
       })
     })
 
+    it('should restore valid UI parts from the top-level message history snapshot', async () => {
+      const uiPart = createHistoricalUIPart()
+      mockFetchConversations.mockResolvedValue(createConversationData())
+      mockFetchChatList.mockResolvedValue({
+        data: [
+          {
+            id: 'msg-with-ui',
+            query: 'What is the weather?',
+            answer: '',
+            message_files: [],
+            agent_thoughts: null,
+            extra_contents: [],
+            ui_parts: [uiPart],
+          },
+        ],
+      })
+
+      const { result } = await renderWithClient(() => useChatWithHistory())
+
+      await waitFor(() => {
+        expect(result!.current.appPrevChatTree).toHaveLength(1)
+      })
+      const answerNode = result!.current.appPrevChatTree[0]?.children?.[0]
+      expect(answerNode?.ui_parts).toEqual([uiPart])
+    })
+
     it('should build tree for paused message with human_input extra_content', async () => {
       // Arrange
       const listData = createConversationData({

web/app/components/base/chat/chat-with-history/hooks.tsx

@@ -34,6 +34,7 @@ import {
 import { TransferMethod } from '@/types/app'
 import { addFileInfos, sortAgentSorts } from '../../../tools/utils'
 import { enrichSubmittedHumanInputFormData } from '../chat/answer/human-input-content/submitted-utils'
+import { normalizeUIParts } from '../chat/ui-part/state'
 import {
   buildChatItemTree,
   getProcessedSystemVariablesFromUrlParams,
@@ -112,6 +113,7 @@ function getFormattedChatList(messages: any[]) {
       citation: item.retriever_resources,
       reasoningContent: item.metadata?.reasoning,
       reasoningFinished: true,
+      ui_parts: normalizeUIParts(item.ui_parts ?? item.metadata?.ui_parts),
       message_files: getProcessedFilesFromResponse(
         answerFiles.map((item: any) => ({
           ...item,

web/app/components/base/chat/chat/__tests__/hooks.spec.tsx

@@ -1,5 +1,6 @@
 import type { ChatConfig, ChatItemInTree } from '../../types'
 import type { FileEntity } from '@/app/components/base/file-uploader/types'
+import type { UIPart } from '@/types/a2ui'
 import { act, renderHook } from '@testing-library/react'
 import { InputVarType, WorkflowRunningStatus } from '@/app/components/workflow/types'
 import { useParams, usePathname } from '@/next/navigation'
@@ -49,12 +50,39 @@ const createAbortControllerMock = () => {
   vi.spyOn(controller, 'abort')
   return controller
 }
+
+function createTestUIPart(sequence = 1): UIPart {
+  return {
+    part_id: 'weather',
+    sequence,
+    protocol: 'a2ui',
+    protocol_version: 'v0.9.1',
+    messages: [
+      {
+        version: 'v0.9.1',
+        createSurface: {
+          surfaceId: 'forecast',
+          catalogId: 'https://dify.ai/a2ui/catalog/v1',
+        },
+      },
+      {
+        version: 'v0.9.1',
+        updateComponents: {
+          surfaceId: 'forecast',
+          components: [{ id: 'root', component: 'Text', text: `Forecast ${sequence}` }],
+        },
+      },
+    ],
+  }
+}
+
 type HookCallbacks = {
   getAbortController: (abortController: AbortController) => void
   onCompleted: (hasError?: boolean, errorMessage?: string) => Promise<void> | void
   onData: (message: string, isFirstMessage: boolean, moreInfo: Record<string, unknown>) => void
   onThought: (thought: Record<string, unknown>) => void
   onFile: (file: Record<string, unknown>) => void
+  onUIPart: (event: { id: string; part: UIPart }) => void
   onMessageEnd: (messageEnd: Record<string, unknown>) => void
   onMessageReplace: (messageReplace: Record<string, unknown>) => void
   onError: (...args: unknown[]) => void
@@ -657,6 +685,7 @@ describe('useChat', () => {
 
     it('should fetch conversation messages and suggested questions onCompleted', async () => {
       let callbacks: HookCallbacks
+      const historyUIPart = createTestUIPart(2)
 
       vi.mocked(ssePost).mockImplementation(async (url, params, options) => {
         callbacks = options as HookCallbacks
@@ -674,6 +703,7 @@ describe('useChat', () => {
             message_tokens: 5,
             provider_response_latency: 0.5,
             workflow_run_id: 'workflow-run-from-history',
+            ui_parts: [historyUIPart],
             inputs: {},
             query: 'hi',
           },
@@ -717,6 +747,7 @@ describe('useChat', () => {
 
       const updatedResponse = result.current.chatList[1]
       expect(updatedResponse!.content).toBe('Updated answer from history') // Fetched from mock
+      expect(updatedResponse!.ui_parts).toEqual([historyUIPart])
       expect(result.current.suggestedQuestions).toEqual(['Suggested 1', 'Suggested 2'])
     })
 
@@ -2053,6 +2084,40 @@ describe('useChat', () => {
         result.current.handleSwitchSibling('a-deep', { isPublicAPI: true })
       })
     })
+
+    it('preserves streamed UI parts when message_end omits the final snapshot', () => {
+      let callbacks: HookCallbacks
+      vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => {
+        callbacks = options as HookCallbacks
+      })
+      const streamedUIPart = createTestUIPart()
+      const prevChatTree = [
+        {
+          id: 'q-ui',
+          content: 'query',
+          isAnswer: false,
+          children: [
+            {
+              id: 'm-ui',
+              content: 'initial',
+              isAnswer: true,
+              siblingIndex: 0,
+              ui_parts: [streamedUIPart],
+            },
+          ],
+        },
+      ]
+      const { result } = renderHook(() =>
+        useChat(undefined, undefined, prevChatTree as ChatItemInTree[]),
+      )
+
+      act(() => {
+        result.current.handleResume('m-ui', 'wr-ui', { isPublicAPI: true })
+        callbacks.onMessageEnd({ id: 'm-ui', metadata: {} })
+      })
+
+      expect(result.current.chatList[1]!.ui_parts).toEqual([streamedUIPart])
+    })
   })
 
   describe('Uncovered edge cases', () => {
@@ -2116,6 +2181,27 @@ describe('useChat', () => {
       expect(lastResponse!.citation).toEqual([])
     })
 
+    it('uses message_end UI parts as the final stream snapshot', () => {
+      let callbacks: HookCallbacks
+      vi.mocked(ssePost).mockImplementation(async (_url, _params, options) => {
+        callbacks = options as HookCallbacks
+      })
+      const streamedUIPart = createTestUIPart()
+      const finalUIPart = createTestUIPart(2)
+      const { result } = renderHook(() => useChat())
+
+      act(() => {
+        result.current.handleSend('url', { query: 'weather' }, {})
+        callbacks.onUIPart({ id: 'm-ui', part: streamedUIPart })
+        callbacks.onMessageEnd({
+          id: 'm-ui',
+          metadata: { ui_parts: [finalUIPart] },
+        })
+      })
+
+      expect(result.current.chatList[1]!.ui_parts).toEqual([finalUIPart])
+    })
+
     it('should handle iteration and loop tracing edge cases (lazy arrays, node finish index -1)', () => {
       let callbacks: HookCallbacks
       vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => {

web/app/components/base/chat/chat/answer/__tests__/index.spec.tsx

@@ -104,6 +104,54 @@ describe('Answer Component', () => {
       expect(screen.getByTestId('custom-agent-content')).toHaveTextContent('streamed answer')
     })
 
+    it('should render a UI part when the answer has no text', () => {
+      render(
+        <Answer
+          {...defaultProps}
+          item={
+            {
+              id: 'msg-with-ui',
+              content: '',
+              isAnswer: true,
+              ui_parts: [
+                {
+                  part_id: 'clock',
+                  sequence: 1,
+                  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: 'Card',
+                            children: [],
+                            title: '14:30',
+                          },
+                        ],
+                      },
+                    },
+                  ],
+                },
+              ],
+            } as unknown as ChatItem
+          }
+        />,
+      )
+
+      expect(screen.getByRole('heading', { name: '14:30' })).toBeInTheDocument()
+    })
+
     it('should render file lists', () => {
       render(
         <Answer

web/app/components/base/chat/chat/answer/index.tsx

@@ -1,4 +1,4 @@
-import type { FC, ReactNode } from 'react'
+import type { ReactNode } from 'react'
 import type { ChatConfig, ChatItem } from '../../types'
 import type { HumanInputFormSubmitData } from './human-input-content/type'
 import type { AppData } from '@/models/share'
@@ -9,6 +9,7 @@ import { EditTitle } from '@/app/components/app/annotation/edit-annotation-modal
 import AnswerIcon from '@/app/components/base/answer-icon'
 import Citation from '@/app/components/base/chat/chat/citation'
 import LoadingAnim from '@/app/components/base/chat/chat/loading-anim'
+import { UIPartList } from '@/app/components/base/chat/chat/ui-part/renderer'
 import { FileList } from '@/app/components/base/file-uploader'
 import ContentSwitch from '../content-switch'
 import { useChatContext } from '../context'
@@ -43,7 +44,7 @@ type AnswerProps = {
   }) => ReactNode
   onHumanInputFormSubmit?: (formToken: string, formData: HumanInputFormSubmitData) => Promise<void>
 }
-const Answer: FC<AnswerProps> = ({
+function Answer({
   item,
   question,
   index,
@@ -59,7 +60,7 @@ const Answer: FC<AnswerProps> = ({
   hideAvatar,
   renderAgentContent,
   onHumanInputFormSubmit,
-}) => {
+}: AnswerProps) {
   const { t } = useTranslation()
   const {
     content,
@@ -72,10 +73,12 @@ const Answer: FC<AnswerProps> = ({
     message_files,
     humanInputFormDataList,
     humanInputFilledFormDataList,
+    ui_parts,
   } = item
   const hasAgentThoughts = !!agent_thoughts?.length
   const hasAgentResponseParts = !!item.agent_response_parts?.length
   const hasAgentContent = hasAgentThoughts || hasAgentResponseParts
+  const hasUIParts = !!ui_parts?.length
   const hasHumanInputs = !!humanInputFormDataList?.length || !!humanInputFilledFormDataList?.length
   // Truthy only when there is real reasoning text. Rehydrated messages carry an empty
   // `{}` (the field is always persisted), and `!!{}` would otherwise be truthy.
@@ -183,7 +186,7 @@ const Answer: FC<AnswerProps> = ({
                 'relative inline-block w-full max-w-full rounded-2xl bg-chat-bubble-bg px-4 py-3 body-lg-regular text-text-primary',
               )}
             >
-              {!responding && contentIsEmpty && !hasAgentContent && (
+              {!responding && contentIsEmpty && !hasAgentContent && !hasUIParts && (
                 <Operation
                   hasWorkflowProcess={!!workflowProcess}
                   maxSize={containerWidth - humanInputFormContainerWidth - 4}
@@ -222,7 +225,8 @@ const Answer: FC<AnswerProps> = ({
                 item.siblingCount > 1 &&
                 !responding &&
                 contentIsEmpty &&
-                !hasAgentContent && (
+                !hasAgentContent &&
+                !hasUIParts && (
                   <ContentSwitch
                     count={item.siblingCount}
                     currentIndex={item.siblingIndex}
@@ -236,75 +240,84 @@ const Answer: FC<AnswerProps> = ({
         )}
 
         {/* Block 2: Response Content (when human inputs exist) */}
-        {hasHumanInputs && (responding || !contentIsEmpty || hasAgentContent || hasReasoning) && (
-          <div className={cn('group relative mt-2 pr-10', chatAnswerContainerInner)}>
-            <div className="absolute -top-2 left-6 h-3 w-0.5 bg-chat-answer-human-input-form-divider-bg" />
-            <div
-              ref={contentRef}
-              className="relative inline-block w-full max-w-full rounded-2xl bg-chat-bubble-bg px-4 py-3 body-lg-regular text-text-primary"
-            >
-              {!responding && (
-                <Operation
-                  hasWorkflowProcess={!!workflowProcess}
-                  maxSize={containerWidth - contentWidth - 4}
-                  contentWidth={contentWidth}
-                  item={item}
-                  question={question}
-                  index={index}
-                  showPromptLog={showPromptLog}
-                  noChatInput={noChatInput}
-                />
-              )}
-              {hasReasoning && (
-                <ReasoningPanel content={item.reasoningContent ?? {}} done={reasoningDone} />
-              )}
-              {responding && contentIsEmpty && !hasAgentContent && !hasReasoning && (
-                <div className="flex h-5 w-6 items-center justify-center">
-                  <LoadingAnim type="text" />
-                </div>
-              )}
-              {!contentIsEmpty && !hasAgentContent && <BasicContent item={item} />}
-              {hasAgentContent && agentContentNode}
-              {!!allFiles?.length && (
-                <FileList
-                  className="my-1"
-                  files={allFiles}
-                  showDeleteAction={false}
-                  showDownloadAction
-                  canPreview
-                />
-              )}
-              {!!message_files?.length && (
-                <FileList
-                  className="my-1"
-                  files={message_files}
-                  showDeleteAction={false}
-                  showDownloadAction
-                  canPreview
-                />
-              )}
-              {annotation?.id && annotation.authorName && (
-                <EditTitle
-                  className="mt-1"
-                  title={t(($) => $.editBy, { ns: 'appAnnotation', author: annotation.authorName })}
-                />
-              )}
-              <SuggestedQuestions item={item} />
-              {!!citation?.length && !responding && (
-                <Citation data={citation} showHitInfo={config?.supportCitationHitInfo} />
-              )}
-              {typeof item.siblingCount === 'number' && item.siblingCount > 1 && (
-                <ContentSwitch
-                  count={item.siblingCount}
-                  currentIndex={item.siblingIndex}
-                  prevDisabled={!item.prevSibling}
-                  nextDisabled={!item.nextSibling}
-                  switchSibling={handleSwitchSibling}
-                />
-              )}
+        {hasHumanInputs &&
+          (responding || !contentIsEmpty || hasAgentContent || hasReasoning || hasUIParts) && (
+            <div className={cn('group relative mt-2 pr-10', chatAnswerContainerInner)}>
+              <div className="absolute -top-2 left-6 h-3 w-0.5 bg-chat-answer-human-input-form-divider-bg" />
+              <div
+                ref={contentRef}
+                className="relative inline-block w-full max-w-full rounded-2xl bg-chat-bubble-bg px-4 py-3 body-lg-regular text-text-primary"
+              >
+                {!responding && (
+                  <Operation
+                    hasWorkflowProcess={!!workflowProcess}
+                    maxSize={containerWidth - contentWidth - 4}
+                    contentWidth={contentWidth}
+                    item={item}
+                    question={question}
+                    index={index}
+                    showPromptLog={showPromptLog}
+                    noChatInput={noChatInput}
+                  />
+                )}
+                {hasReasoning && (
+                  <ReasoningPanel content={item.reasoningContent ?? {}} done={reasoningDone} />
+                )}
+                {responding &&
+                  contentIsEmpty &&
+                  !hasAgentContent &&
+                  !hasReasoning &&
+                  !hasUIParts && (
+                    <div className="flex h-5 w-6 items-center justify-center">
+                      <LoadingAnim type="text" />
+                    </div>
+                  )}
+                {!contentIsEmpty && !hasAgentContent && <BasicContent item={item} />}
+                {hasAgentContent && agentContentNode}
+                {hasUIParts && <UIPartList parts={ui_parts} />}
+                {!!allFiles?.length && (
+                  <FileList
+                    className="my-1"
+                    files={allFiles}
+                    showDeleteAction={false}
+                    showDownloadAction
+                    canPreview
+                  />
+                )}
+                {!!message_files?.length && (
+                  <FileList
+                    className="my-1"
+                    files={message_files}
+                    showDeleteAction={false}
+                    showDownloadAction
+                    canPreview
+                  />
+                )}
+                {annotation?.id && annotation.authorName && (
+                  <EditTitle
+                    className="mt-1"
+                    title={t(($) => $.editBy, {
+                      ns: 'appAnnotation',
+                      author: annotation.authorName,
+                    })}
+                  />
+                )}
+                <SuggestedQuestions item={item} />
+                {!!citation?.length && !responding && (
+                  <Citation data={citation} showHitInfo={config?.supportCitationHitInfo} />
+                )}
+                {typeof item.siblingCount === 'number' && item.siblingCount > 1 && (
+                  <ContentSwitch
+                    count={item.siblingCount}
+                    currentIndex={item.siblingIndex}
+                    prevDisabled={!item.prevSibling}
+                    nextDisabled={!item.nextSibling}
+                    switchSibling={handleSwitchSibling}
+                  />
+                )}
+              </div>
             </div>
-          </div>
-        )}
+          )}
 
         {/* Original single block layout (when no human inputs) */}
         {!hasHumanInputs && (
@@ -316,7 +329,7 @@ const Answer: FC<AnswerProps> = ({
               ref={contentRef}
               className={cn(
                 'relative inline-block max-w-full rounded-2xl bg-chat-bubble-bg px-4 py-3 body-lg-regular text-text-primary',
-                workflowProcess && 'w-full',
+                (workflowProcess || hasUIParts) && 'w-full',
               )}
             >
               {!responding && (
@@ -345,13 +358,14 @@ const Answer: FC<AnswerProps> = ({
               {hasReasoning && (
                 <ReasoningPanel content={item.reasoningContent ?? {}} done={reasoningDone} />
               )}
-              {responding && contentIsEmpty && !hasAgentContent && !hasReasoning && (
+              {responding && contentIsEmpty && !hasAgentContent && !hasReasoning && !hasUIParts && (
                 <div className="flex h-5 w-6 items-center justify-center">
                   <LoadingAnim type="text" />
                 </div>
               )}
               {!contentIsEmpty && !hasAgentContent && <BasicContent item={item} />}
               {hasAgentContent && agentContentNode}
+              {hasUIParts && <UIPartList parts={ui_parts} />}
               {!!allFiles?.length && (
                 <FileList
                   className="my-1"

web/app/components/base/chat/chat/hooks.ts

@@ -15,6 +15,10 @@ import { useTranslation } from 'react-i18next'
 import { v4 as uuidV4 } from 'uuid'
 import { AudioPlayerManager } from '@/app/components/base/audio-btn/audio.player.manager'
 import { enrichSubmittedHumanInputFormData } from '@/app/components/base/chat/chat/answer/human-input-content/submitted-utils'
+import {
+  reconcileHistoryUIParts,
+  upsertUIPart,
+} from '@/app/components/base/chat/chat/ui-part/state'
 import {
   getProcessedFiles,
   getProcessedFilesFromResponse,
@@ -43,7 +47,9 @@ type HistoryConversationMessage = {
   retriever_resources?: ChatItem['citation']
   metadata?: {
     reasoning?: ChatItem['reasoningContent']
+    ui_parts?: unknown
   }
+  ui_parts?: unknown
   created_at?: number
   answer_tokens?: number
   message_tokens?: number
@@ -550,8 +556,18 @@ export const useChat = (
             }
           })
         },
+        onUIPart({ id, part }) {
+          updateChatTreeNode(messageId, (responseItem) => {
+            if (id) responseItem.id = id
+            responseItem.ui_parts = upsertUIPart(responseItem.ui_parts ?? [], part)
+          })
+        },
         onMessageEnd: (messageEnd) => {
           updateChatTreeNode(messageId, (responseItem) => {
+            responseItem.ui_parts = reconcileHistoryUIParts(
+              responseItem.ui_parts,
+              messageEnd.metadata?.ui_parts,
+            )
             if (messageEnd.metadata?.annotation_reply) {
               responseItem.annotation = {
                 id: messageEnd.metadata.annotation_reply.id,
@@ -1045,6 +1061,10 @@ export const useChat = (
                 citation: newResponseItem.retriever_resources,
                 reasoningContent: newResponseItem.metadata?.reasoning,
                 reasoningFinished: true,
+                ui_parts: reconcileHistoryUIParts(
+                  responseItem.ui_parts,
+                  newResponseItem.ui_parts ?? newResponseItem.metadata?.ui_parts,
+                ),
                 message_files: historyAnswerFiles,
                 allFiles: undefined,
                 workflowProcess: undefined,
@@ -1198,7 +1218,26 @@ export const useChat = (
             parentId: data.parent_message_id,
           })
         },
+        onUIPart({ id, part }) {
+          if (id && !hasSetResponseId) {
+            questionItem.id = `question-${id}`
+            responseItem.id = id
+            responseItem.parentMessageId = questionItem.id
+            hasSetResponseId = true
+          }
+          responseItem.ui_parts = upsertUIPart(responseItem.ui_parts ?? [], part)
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
+          })
+        },
         onMessageEnd: (messageEnd) => {
+          responseItem.ui_parts = reconcileHistoryUIParts(
+            responseItem.ui_parts,
+            messageEnd.metadata?.ui_parts,
+          )
           if (messageEnd.metadata?.annotation_reply) {
             responseItem.id = messageEnd.id
             responseItem.annotation = {

web/app/components/base/chat/chat/type.ts

@@ -2,6 +2,7 @@ import type { FileEntity } from '@/app/components/base/file-uploader/types'
 import type { TypeWithI18N } from '@/app/components/header/account-setting/model-provider-page/declarations'
 import type { InputVarType } from '@/app/components/workflow/types'
 import type { Annotation, MessageRating } from '@/models/log'
+import type { UIPart } from '@/types/a2ui'
 import type { FileResponse, HumanInputFilledFormData, HumanInputFormData } from '@/types/workflow'
 
 type MessageMore = {
@@ -115,6 +116,7 @@ export type IChatItem = {
   log?: { role: string; text: string; files?: FileEntity[] }[]
   agent_thoughts?: ThoughtItem[]
   agent_response_parts?: AgentResponsePart[]
+  ui_parts?: UIPart[]
   // for LLM reasoning (chain-of-thought) in "separated" mode, keyed by LLM node id
   reasoningContent?: Record<string, string>
   reasoningFinished?: boolean
@@ -136,6 +138,7 @@ export type IChatItem = {
 
 export type Metadata = {
   retriever_resources?: CitationItem[]
+  ui_parts?: unknown
   annotation_reply: {
     id: string
     account: {

web/app/components/base/chat/chat/ui-part/README.md

@@ -0,0 +1,11 @@
+# Chat UI Part
+
+Renders validated, read-only A2UI v0.9.1 tool output inside an assistant chat message.
+
+## Internal Modules
+
+None.
+
+## External Modules
+
+- `types/a2ui`

web/app/components/base/chat/chat/ui-part/__tests__/renderer.spec.tsx

@@ -0,0 +1,180 @@
+import type { UIPart } from '@/types/a2ui'
+import { render, screen } from '@testing-library/react'
+import { describe, expect, it } from 'vitest'
+import { UIPartList } from '../renderer'
+
+function createPart(
+  messages: UIPart['messages'],
+  fallback = 'Unable to display this widget.',
+): UIPart {
+  return {
+    part_id: 'weather',
+    sequence: 1,
+    protocol: 'a2ui',
+    protocol_version: 'v0.9.1',
+    messages,
+    fallback,
+  }
+}
+
+describe('UIPartList', () => {
+  it('renders a root-adjacent read-only surface with data bindings', () => {
+    render(
+      <UIPartList
+        parts={[
+          createPart([
+            {
+              version: 'v0.9.1',
+              createSurface: {
+                surfaceId: 'forecast',
+                catalogId: 'https://dify.ai/a2ui/catalog/v1',
+              },
+            },
+            {
+              version: 'v0.9.1',
+              updateDataModel: {
+                surfaceId: 'forecast',
+                value: {
+                  city: 'Shanghai',
+                  temperature: 31,
+                },
+              },
+            },
+            {
+              version: 'v0.9.1',
+              updateComponents: {
+                surfaceId: 'forecast',
+                components: [
+                  { id: 'root', component: 'Card', children: ['city', 'temperature'] },
+                  {
+                    id: 'city',
+                    component: 'Card',
+                    children: [],
+                    title: { path: '/city' },
+                  },
+                  {
+                    id: 'temperature',
+                    component: 'Metric',
+                    label: 'Temperature',
+                    value: { path: '/temperature' },
+                    unit: '°C',
+                  },
+                ],
+              },
+            },
+          ]),
+        ]}
+      />,
+    )
+
+    expect(screen.getByRole('heading', { name: 'Shanghai' })).toBeInTheDocument()
+    expect(screen.getByText('31')).toBeInTheDocument()
+    expect(screen.getByText('°C')).toBeInTheDocument()
+  })
+
+  it('uses fallback content when the component graph has a cycle', () => {
+    render(
+      <UIPartList
+        parts={[
+          createPart(
+            [
+              {
+                version: 'v0.9.1',
+                createSurface: {
+                  surfaceId: 'cycle',
+                  catalogId: 'https://dify.ai/a2ui/catalog/v1',
+                },
+              },
+              {
+                version: 'v0.9.1',
+                updateComponents: {
+                  surfaceId: 'cycle',
+                  components: [
+                    { id: 'root', component: 'Card', children: ['loop'] },
+                    { id: 'loop', component: 'Column', children: ['root'] },
+                  ],
+                },
+              },
+            ],
+            'Weather summary is unavailable.',
+          ),
+        ]}
+      />,
+    )
+
+    expect(screen.getByText('Weather summary is unavailable.')).toBeInTheDocument()
+  })
+
+  it('fails closed instead of expanding a shared-child component graph', () => {
+    render(
+      <UIPartList
+        parts={[
+          createPart(
+            [
+              {
+                version: 'v0.9.1',
+                createSurface: {
+                  surfaceId: 'shared-child',
+                  catalogId: 'https://dify.ai/a2ui/catalog/v1',
+                },
+              },
+              {
+                version: 'v0.9.1',
+                updateComponents: {
+                  surfaceId: 'shared-child',
+                  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: 'Do not expand me' },
+                  ],
+                },
+              },
+            ],
+            'Widget graph was rejected.',
+          ),
+        ]}
+      />,
+    )
+
+    expect(screen.getByText('Widget graph was rejected.')).toBeInTheDocument()
+    expect(screen.queryByText('Do not expand me')).not.toBeInTheDocument()
+  })
+
+  it('renders Progress with its required accessible label', () => {
+    render(
+      <UIPartList
+        parts={[
+          createPart([
+            {
+              version: 'v0.9.1',
+              createSurface: {
+                surfaceId: 'progress',
+                catalogId: 'https://dify.ai/a2ui/catalog/v1',
+              },
+            },
+            {
+              version: 'v0.9.1',
+              updateComponents: {
+                surfaceId: 'progress',
+                components: [
+                  {
+                    id: 'root',
+                    component: 'Progress',
+                    value: 75,
+                    label: 'Forecast loaded',
+                  },
+                ],
+              },
+            },
+          ]),
+        ]}
+      />,
+    )
+
+    expect(screen.getByRole('progressbar', { name: 'Forecast loaded' })).toHaveAttribute(
+      'aria-valuenow',
+      '75',
+    )
+  })
+})

web/app/components/base/chat/chat/ui-part/__tests__/state.spec.ts

@@ -0,0 +1,786 @@
+import type { A2UIComponent, UIPart } from '@/types/a2ui'
+import { describe, expect, it } from 'vitest'
+import {
+  buildUISurfaces,
+  normalizeUIParts,
+  parseUIPart,
+  parseUIPartEnvelope,
+  reconcileHistoryUIParts,
+  resolveJSONPointer,
+  upsertUIPart,
+} from '../state'
+
+function createPart(overrides: Partial<UIPart> = {}): UIPart {
+  return {
+    part_id: 'weather',
+    sequence: 1,
+    protocol: 'a2ui',
+    protocol_version: 'v0.9.1',
+    messages: [
+      {
+        version: 'v0.9.1',
+        createSurface: {
+          surfaceId: 'forecast',
+          catalogId: 'https://dify.ai/a2ui/catalog/v1',
+        },
+      },
+      {
+        version: 'v0.9.1',
+        updateDataModel: {
+          surfaceId: 'forecast',
+          value: {
+            city: 'Shanghai',
+            temperature: 31,
+          },
+        },
+      },
+      {
+        version: 'v0.9.1',
+        updateComponents: {
+          surfaceId: 'forecast',
+          components: [
+            { id: 'root', component: 'Card', children: ['city', 'temperature'] },
+            { id: 'city', component: 'Text', text: { path: '/city' } },
+            {
+              id: 'temperature',
+              component: 'Metric',
+              label: 'Temperature',
+              value: { path: '/temperature' },
+              unit: '°C',
+            },
+          ],
+        },
+      },
+    ],
+    fallback: 'Weather is unavailable.',
+    ...overrides,
+  }
+}
+
+function createLargePart(partId: string): UIPart {
+  const largeModel = Object.fromEntries(
+    Array.from({ length: 20 }, (_, index) => [`value-${index}`, 'x'.repeat(4_096)]),
+  )
+  return createPart({
+    part_id: partId,
+    messages: [
+      {
+        version: 'v0.9.1',
+        createSurface: {
+          surfaceId: 'large',
+          catalogId: 'https://dify.ai/a2ui/catalog/v1',
+        },
+      },
+      {
+        version: 'v0.9.1',
+        updateDataModel: {
+          surfaceId: 'large',
+          value: largeModel,
+        },
+      },
+      {
+        version: 'v0.9.1',
+        updateComponents: {
+          surfaceId: 'large',
+          components: [{ id: 'root', component: 'Text', text: 'Large widget' }],
+        },
+      },
+    ],
+  })
+}
+
+describe('ui part state', () => {
+  it('keeps insertion order while replacing a part only with a newer sequence', () => {
+    const weather = createPart()
+    const clock = createPart({ part_id: 'clock', sequence: 2 })
+    const initial = upsertUIPart(upsertUIPart([], weather), clock)
+
+    expect(upsertUIPart(initial, createPart({ sequence: 0 }))).toEqual(initial)
+
+    const updated = upsertUIPart(initial, createPart({ sequence: 3, fallback: 'updated' }))
+    expect(updated.map((part) => part.part_id)).toEqual(['weather', 'clock'])
+    expect(updated[0]?.sequence).toBe(3)
+    expect(updated[0]?.fallback).toBe('updated')
+  })
+
+  it('keeps existing live parts when a seventeenth distinct part exceeds the collection limit', () => {
+    const liveParts = Array.from({ length: 16 }, (_, index) =>
+      createPart({ part_id: `part-${index}` }),
+    )
+    const overflowPart = createPart({ part_id: 'part-16' })
+
+    expect(upsertUIPart(liveParts, overflowPart)).toBe(liveParts)
+    expect(normalizeUIParts([...liveParts, overflowPart])).toEqual([])
+    expect(reconcileHistoryUIParts(liveParts, [...liveParts, overflowPart])).toBe(liveParts)
+  })
+
+  it('keeps existing live parts when history exceeds the 512 KiB collection budget', () => {
+    const liveParts = [createPart()]
+    const oversizedHistory = Array.from({ length: 7 }, (_, index) =>
+      createLargePart(`large-${index}`),
+    )
+    const withinBudget = oversizedHistory.slice(0, 6)
+
+    expect(upsertUIPart(withinBudget, oversizedHistory[6])).toBe(withinBudget)
+    expect(normalizeUIParts(oversizedHistory)).toEqual([])
+    expect(reconcileHistoryUIParts(liveParts, oversizedHistory)).toBe(liveParts)
+  })
+
+  it('limits history normalization to sixty-four candidates even for repeated part ids', () => {
+    const liveParts = [createPart({ sequence: 100 })]
+    const boundedHistory = Array.from({ length: 64 }, (_, index) =>
+      createPart({ sequence: index + 1 }),
+    )
+    const oversizedHistory = [
+      ...boundedHistory,
+      createPart({ sequence: boundedHistory.length + 1 }),
+    ]
+
+    expect(normalizeUIParts(boundedHistory)).toEqual([createPart({ sequence: 64 })])
+    expect(normalizeUIParts(oversizedHistory)).toEqual([])
+    expect(reconcileHistoryUIParts(liveParts, oversizedHistory)).toBe(liveParts)
+  })
+
+  it('retains live UI when an older history response omits ui_parts', () => {
+    const liveParts = [createPart()]
+
+    expect(reconcileHistoryUIParts(liveParts, undefined)).toBe(liveParts)
+    expect(reconcileHistoryUIParts(liveParts, [])).toEqual([])
+  })
+
+  it('rejects envelopes outside the fixed catalog and component schema', () => {
+    const unknownComponent = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [
+              {
+                id: 'root',
+                component: 'Button',
+                action: 'buy',
+              } as never,
+            ],
+          },
+        },
+      ],
+    })
+    const styledComponent = createPart({
+      part_id: 'styled',
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [
+              {
+                id: 'root',
+                component: 'Text',
+                text: 'unsafe',
+                className: 'fixed inset-0',
+              } as never,
+            ],
+          },
+        },
+      ],
+    })
+    const imageComponent = createPart({
+      part_id: 'image',
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [
+              {
+                id: 'root',
+                component: 'Image',
+                src: 'https://example.com/weather.png',
+              } as never,
+            ],
+          },
+        },
+      ],
+    })
+
+    expect(normalizeUIParts([unknownComponent, styledComponent, imageComponent])).toEqual([])
+  })
+
+  it('requires an accessible label for Progress components', () => {
+    const unlabeledProgress = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [
+              {
+                id: 'root',
+                component: 'Progress',
+                value: 50,
+              } as never,
+            ],
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(unlabeledProgress)).toBeUndefined()
+  })
+
+  it('rejects duplicate component ids within a single update', () => {
+    const duplicateComponents = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [
+              { id: 'root', component: 'Text', text: 'First' },
+              { id: 'root', component: 'Text', text: 'Second' },
+            ],
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(duplicateComponents)).toBeUndefined()
+  })
+
+  it('applies the single-part byte limit to the API payload without stream wrapper fields', () => {
+    const repeatedModelUpdate: UIPart['messages'][number] = {
+      version: 'v0.9.1',
+      updateDataModel: {
+        surfaceId: 'forecast',
+        value: { payload: 'x'.repeat(4_096) },
+      },
+    }
+    const messages: UIPart['messages'] = [
+      {
+        version: 'v0.9.1',
+        createSurface: {
+          surfaceId: 'forecast',
+          catalogId: 'https://dify.ai/a2ui/catalog/v1',
+        },
+      },
+      ...Array.from({ length: 31 }, () => repeatedModelUpdate),
+      {
+        version: 'v0.9.1',
+        updateComponents: {
+          surfaceId: 'forecast',
+          components: [{ id: 'root', component: 'Text', text: 'Widget' }],
+        },
+      },
+    ]
+    const payloadWithoutFallback = {
+      protocol: 'a2ui' as const,
+      protocol_version: 'v0.9.1' as const,
+      messages,
+    }
+    const payloadWithEmptyFallbackBytes = new TextEncoder().encode(
+      JSON.stringify({ ...payloadWithoutFallback, fallback: '' }),
+    ).byteLength
+    const targetPayloadBytes = 128 * 1_024 - 100
+    const fallback = 'x'.repeat(targetPayloadBytes - payloadWithEmptyFallbackBytes)
+    const part: UIPart = {
+      part_id: 'p'.repeat(512),
+      sequence: Number.MAX_SAFE_INTEGER,
+      ...payloadWithoutFallback,
+      fallback,
+    }
+
+    expect(
+      new TextEncoder().encode(
+        JSON.stringify({
+          protocol: part.protocol,
+          protocol_version: part.protocol_version,
+          messages: part.messages,
+          fallback: part.fallback,
+        }),
+      ).byteLength,
+    ).toBe(targetPayloadBytes)
+    expect(new TextEncoder().encode(JSON.stringify(part)).byteLength).toBeGreaterThan(128 * 1_024)
+    expect(parseUIPart(part)).toEqual(part)
+  })
+
+  it.each<{
+    name: string
+    messages: UIPart['messages']
+  }>([
+    {
+      name: 'an update before createSurface',
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+      ],
+    },
+    {
+      name: 'a repeated createSurface',
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+      ],
+    },
+    {
+      name: 'messages targeting multiple surfaces',
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'other',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+      ],
+    },
+    {
+      name: 'deleteSurface before the final message',
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          deleteSurface: {
+            surfaceId: 'forecast',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+      ],
+    },
+    {
+      name: 'more than one deleteSurface',
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          deleteSurface: {
+            surfaceId: 'forecast',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          deleteSurface: {
+            surfaceId: 'forecast',
+          },
+        },
+      ],
+    },
+  ])('rejects $name', ({ messages }) => {
+    expect(parseUIPart(createPart({ messages }))).toBeUndefined()
+  })
+
+  it('accepts tool-scoped part ids and enforces the 100-component catalog bound', () => {
+    const components = Array.from({ length: 100 }, (_, index) => ({
+      id: index === 0 ? 'root' : `component-${index}`,
+      component: 'Divider' as const,
+    }))
+    const boundedPart = createPart({
+      part_id: 'call/abc=:weather',
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components,
+          },
+        },
+      ],
+    })
+    const oversizedPart = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [
+              ...components,
+              {
+                id: 'component-100',
+                component: 'Divider',
+              },
+            ],
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(boundedPart)).toEqual(boundedPart)
+    expect(parseUIPartEnvelope(oversizedPart)).toBeUndefined()
+  })
+
+  it('enforces the same 100-item bound on component children', () => {
+    const children = Array.from({ length: 100 }, (_, index) => `child-${index}`)
+    const boundedPart = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Column', children }],
+          },
+        },
+      ],
+    })
+    const oversizedPart = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [
+              {
+                id: 'root',
+                component: 'Column',
+                children: [...children, 'child-100'],
+              },
+            ],
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(boundedPart)).toEqual(boundedPart)
+    expect(parseUIPartEnvelope(oversizedPart)).toBeUndefined()
+  })
+
+  it('reduces surface messages and removes deleted surfaces', () => {
+    const part = createPart({
+      messages: [
+        ...createPart().messages,
+        {
+          version: 'v0.9.1',
+          deleteSurface: {
+            surfaceId: 'forecast',
+          },
+        },
+      ],
+    })
+
+    const result = buildUISurfaces(part.messages)
+    expect(result.error).toBeUndefined()
+    expect(result.surfaces).toEqual([])
+  })
+
+  it.each<{
+    name: string
+    components: A2UIComponent[]
+    error: string
+  }>([
+    {
+      name: 'a duplicate child in one parent',
+      components: [
+        { id: 'root', component: 'Column', children: ['value', 'value'] },
+        { id: 'value', component: 'Text', text: '31°C' },
+      ],
+      error: 'duplicate_child',
+    },
+    {
+      name: 'a child shared by multiple parents',
+      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: '31°C' },
+      ],
+      error: 'multiple_parents',
+    },
+    {
+      name: 'root referenced as a child',
+      components: [
+        { id: 'root', component: 'Column', children: ['branch'] },
+        { id: 'branch', component: 'Column', children: ['root'] },
+      ],
+      error: 'root_as_child',
+    },
+    {
+      name: 'an unreachable component',
+      components: [
+        { id: 'root', component: 'Column', children: [] },
+        { id: 'orphan', component: 'Text', text: 'hidden' },
+      ],
+      error: 'unreachable_component',
+    },
+  ])('rejects $name while reducing a surface', ({ components, error }) => {
+    const part = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'unsafe',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'unsafe',
+            components,
+          },
+        },
+      ],
+    })
+
+    expect(buildUISurfaces(part.messages)).toEqual({
+      surfaces: [],
+      error,
+    })
+    expect(normalizeUIParts([part])).toEqual([])
+  })
+
+  it('immutably upserts nested data model values through JSON Pointer paths', () => {
+    const originalModel = {
+      current: {
+        temperature: 30,
+      },
+    }
+    const part = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateDataModel: {
+            surfaceId: 'forecast',
+            value: originalModel,
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateDataModel: {
+            surfaceId: 'forecast',
+            path: '/current/temperature',
+            value: 31,
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+      ],
+    })
+
+    const result = buildUISurfaces(part.messages)
+    expect(result.surfaces[0]?.dataModel).toEqual({
+      current: {
+        temperature: 31,
+      },
+    })
+    expect(originalModel.current.temperature).toBe(30)
+  })
+
+  it('rejects JSON Pointer paths deeper than sixteen segments', () => {
+    const path = `/${Array.from({ length: 17 }, (_, index) => `level-${index}`).join('/')}`
+    const part = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          updateDataModel: {
+            surfaceId: 'forecast',
+            path,
+            value: 'too deep',
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(part)).toBeUndefined()
+  })
+
+  it('rejects consecutive patches that make the final model exceed the depth limit', () => {
+    const firstPath = `/${Array.from({ length: 15 }, (_, index) => `level-${index}`).join('/')}`
+    const part = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateDataModel: {
+            surfaceId: 'forecast',
+            path: firstPath,
+            value: {},
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateDataModel: {
+            surfaceId: 'forecast',
+            path: `${firstPath}/level-15`,
+            value: { overflow: true },
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(part)).toEqual(part)
+    expect(buildUISurfaces(part.messages)).toEqual({
+      surfaces: [],
+      error: 'invalid_data_model_update',
+    })
+  })
+
+  it('rejects consecutive patches that make the final model exceed the node limit', () => {
+    const updates: UIPart['messages'] = Array.from({ length: 50 }, (_, index) => ({
+      version: 'v0.9.1',
+      updateDataModel: {
+        surfaceId: 'forecast',
+        path: `/bucket-${index}`,
+        value: Array.from({ length: 40 }, () => 0),
+      },
+    }))
+    const part = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        ...updates,
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(part)).toEqual(part)
+    expect(buildUISurfaces(part.messages)).toEqual({
+      surfaces: [],
+      error: 'invalid_data_model_update',
+    })
+  })
+
+  it('rejects a nested index-1000 patch that would create sparse arrays', () => {
+    const part = createPart({
+      messages: [
+        {
+          version: 'v0.9.1',
+          createSurface: {
+            surfaceId: 'forecast',
+            catalogId: 'https://dify.ai/a2ui/catalog/v1',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateDataModel: {
+            surfaceId: 'forecast',
+            value: { items: [] },
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateDataModel: {
+            surfaceId: 'forecast',
+            path: '/items/1000/1000',
+            value: 'amplified',
+          },
+        },
+        {
+          version: 'v0.9.1',
+          updateComponents: {
+            surfaceId: 'forecast',
+            components: [{ id: 'root', component: 'Text', text: 'Forecast' }],
+          },
+        },
+      ],
+    })
+
+    expect(parseUIPartEnvelope(part)).toEqual(part)
+    expect(buildUISurfaces(part.messages)).toEqual({
+      surfaces: [],
+      error: 'invalid_data_model_update',
+    })
+  })
+
+  it('resolves escaped JSON Pointer segments without traversing unsafe keys', () => {
+    const model = {
+      'weather/current': {
+        '~value': 31,
+      },
+    }
+
+    expect(resolveJSONPointer(model, '/weather~1current/~0value')).toBe(31)
+    expect(resolveJSONPointer(model, '/__proto__/polluted')).toBeUndefined()
+    expect(resolveJSONPointer(model, 'weather/current')).toBeUndefined()
+  })
+})

web/app/components/base/chat/chat/ui-part/renderer.tsx

@@ -0,0 +1,292 @@
+import type { ReactNode } from 'react'
+import type {
+  A2UIBinding,
+  A2UIComponent,
+  A2UIDynamic,
+  JSONPrimitive,
+  JSONValue,
+  UIPart,
+} from '@/types/a2ui'
+import {
+  buildUISurfaces,
+  getSurfaceGraphError,
+  parseUIPartEnvelope,
+  resolveJSONPointer,
+} from './state'
+
+const ICON_CLASSES: Record<string, string> = {
+  calendar: 'i-ri-calendar-line',
+  clock: 'i-ri-time-line',
+  cloud: 'i-ri-cloudy-2-line',
+  location: 'i-ri-map-pin-2-line',
+  rain: 'i-ri-rainy-line',
+  snow: 'i-ri-snowy-line',
+  sun: 'i-ri-sun-line',
+  thermometer: 'i-ri-temp-hot-line',
+  wind: 'i-ri-windy-line',
+}
+
+const GAP_CLASSES = {
+  small: 'gap-1',
+  medium: 'gap-2',
+  large: 'gap-4',
+} as const
+
+const ALIGN_CLASSES = {
+  start: 'items-start',
+  center: 'items-center',
+  end: 'items-end',
+} as const
+
+const BADGE_CLASSES = {
+  neutral: 'bg-background-section text-text-secondary',
+  info: 'bg-state-info-hover text-state-info-solid',
+  success: 'bg-state-success-hover text-state-success-solid',
+  warning: 'bg-state-warning-hover text-state-warning-solid',
+  critical: 'bg-state-destructive-hover text-state-destructive-solid',
+} as const
+
+const MAX_RENDERED_COMPONENTS = 100
+
+function isBinding(value: JSONPrimitive | A2UIBinding): value is A2UIBinding {
+  return typeof value === 'object' && value !== null && 'path' in value
+}
+
+function resolveDynamic<T extends JSONPrimitive>(
+  value: A2UIDynamic<T>,
+  model: JSONValue,
+): JSONValue | undefined {
+  return isBinding(value) ? resolveJSONPointer(model, value.path) : value
+}
+
+function resolveString(
+  value: A2UIDynamic<string> | undefined,
+  model: JSONValue,
+): string | undefined {
+  if (value === undefined) return undefined
+  const resolved = resolveDynamic(value, model)
+  return typeof resolved === 'string' ? resolved : undefined
+}
+
+function resolveNumber(
+  value: A2UIDynamic<number> | undefined,
+  model: JSONValue,
+): number | undefined {
+  if (value === undefined) return undefined
+  const resolved = resolveDynamic(value, model)
+  return typeof resolved === 'number' && Number.isFinite(resolved) ? resolved : undefined
+}
+
+function displayPrimitive(value: JSONValue | undefined) {
+  if (
+    value === null ||
+    typeof value === 'string' ||
+    typeof value === 'number' ||
+    typeof value === 'boolean'
+  )
+    return value === null ? '' : String(value)
+
+  return undefined
+}
+
+function formatDateTime(value: string, format: 'date' | 'time' | 'datetime') {
+  const date = new Date(value)
+  if (Number.isNaN(date.getTime())) return value
+
+  const formatOptions: Intl.DateTimeFormatOptions =
+    format === 'date'
+      ? { dateStyle: 'medium' }
+      : format === 'time'
+        ? { timeStyle: 'short' }
+        : { dateStyle: 'medium', timeStyle: 'short' }
+
+  return new Intl.DateTimeFormat(undefined, formatOptions).format(date)
+}
+
+function renderChildren(
+  component: Extract<A2UIComponent, { component: 'Card' | 'Row' | 'Column' }>,
+  renderComponent: (componentId: string) => ReactNode,
+) {
+  return component.children.map((childId) => (
+    <div key={childId} className="min-w-0">
+      {renderComponent(childId)}
+    </div>
+  ))
+}
+
+function Surface({ surface }: { surface: ReturnType<typeof buildUISurfaces>['surfaces'][number] }) {
+  const graphError = getSurfaceGraphError(surface)
+  if (graphError) return null
+  const renderedComponentIds = new Set<string>()
+  let renderFailed = false
+
+  function renderComponent(componentId: string): ReactNode {
+    if (
+      renderFailed ||
+      renderedComponentIds.has(componentId) ||
+      renderedComponentIds.size >= MAX_RENDERED_COMPONENTS
+    ) {
+      renderFailed = true
+      return null
+    }
+    renderedComponentIds.add(componentId)
+
+    const component = surface.components.get(componentId)
+    if (!component) {
+      renderFailed = true
+      return null
+    }
+    const model = surface.dataModel
+
+    switch (component.component) {
+      case 'Card': {
+        const title = resolveString(component.title, model)
+        return (
+          <section className="w-full rounded-xl border border-components-panel-border bg-background-default-subtle p-4 shadow-xs">
+            {title && <h3 className="mb-3 title-sm-semi-bold text-text-primary">{title}</h3>}
+            <div className="flex flex-col gap-3">{renderChildren(component, renderComponent)}</div>
+          </section>
+        )
+      }
+      case 'Row':
+        return (
+          <div
+            className={`flex flex-wrap ${GAP_CLASSES[component.gap ?? 'medium']} ${ALIGN_CLASSES[component.align ?? 'center']}`}
+          >
+            {renderChildren(component, renderComponent)}
+          </div>
+        )
+      case 'Column':
+        return (
+          <div className={`flex flex-col ${GAP_CLASSES[component.gap ?? 'medium']}`}>
+            {renderChildren(component, renderComponent)}
+          </div>
+        )
+      case 'Text': {
+        const text = resolveString(component.text, model)
+        if (text === undefined) return null
+        if (component.variant === 'caption')
+          return <p className="body-xs-regular text-text-tertiary">{text}</p>
+        return <p className="body-md-regular whitespace-pre-wrap text-text-secondary">{text}</p>
+      }
+      case 'Icon': {
+        const iconClass = ICON_CLASSES[component.name]
+        return iconClass ? (
+          <span className={`${iconClass} size-5 text-text-secondary`} aria-hidden="true" />
+        ) : null
+      }
+      case 'Divider':
+        return <hr className="border-divider-subtle" />
+      case 'Badge': {
+        const text = resolveString(component.text, model)
+        if (text === undefined) return null
+        return (
+          <span
+            className={`inline-flex w-fit rounded-md px-2 py-1 body-xs-medium ${BADGE_CLASSES[component.tone ?? 'neutral']}`}
+          >
+            {text}
+          </span>
+        )
+      }
+      case 'Metric': {
+        const label = resolveString(component.label, model)
+        const value = displayPrimitive(resolveDynamic(component.value, model))
+        const unit = resolveString(component.unit, model)
+        if (label === undefined || value === undefined) return null
+        return (
+          <dl>
+            <dt className="body-xs-medium text-text-tertiary">{label}</dt>
+            <dd className="mt-0.5 flex items-baseline gap-1">
+              <span className="title-xl-semi-bold text-text-primary">{value}</span>
+              {unit && <span className="body-sm-medium text-text-secondary">{unit}</span>}
+            </dd>
+          </dl>
+        )
+      }
+      case 'DateTime': {
+        const value = resolveString(component.value, model)
+        if (value === undefined) return null
+        return (
+          <time dateTime={value} className="body-md-medium text-text-primary">
+            {formatDateTime(value, component.format ?? 'datetime')}
+          </time>
+        )
+      }
+      case 'Progress': {
+        const value = resolveNumber(component.value, model)
+        const max = resolveNumber(component.max, model) ?? 100
+        const label = resolveString(component.label, model)
+        if (value === undefined || label === undefined || max <= 0) return null
+        const normalizedValue = Math.min(Math.max(value, 0), max)
+        const percentage = (normalizedValue / max) * 100
+        return (
+          <div className="w-full">
+            <div className="mb-1 body-xs-medium text-text-secondary">{label}</div>
+            <div
+              role="progressbar"
+              aria-label={label}
+              aria-valuemin={0}
+              aria-valuemax={max}
+              aria-valuenow={normalizedValue}
+              className="h-2 overflow-hidden rounded-full bg-background-section"
+            >
+              <div
+                className="h-full rounded-full bg-components-progress-brand-progress"
+                style={{ width: `${percentage}%` }}
+              />
+            </div>
+          </div>
+        )
+      }
+      case 'KeyValue': {
+        const label = resolveString(component.label, model)
+        const value = displayPrimitive(resolveDynamic(component.value, model))
+        if (label === undefined || value === undefined) return null
+        return (
+          <dl className="flex items-baseline justify-between gap-4">
+            <dt className="body-sm-regular text-text-secondary">{label}</dt>
+            <dd className="text-right body-sm-medium text-text-primary">{value}</dd>
+          </dl>
+        )
+      }
+    }
+  }
+
+  const content = renderComponent('root')
+  return renderFailed ? null : <div data-surface-id={surface.surfaceId}>{content}</div>
+}
+
+function Part({ part }: { part: UIPart }) {
+  const validPart = parseUIPartEnvelope(part)
+  if (!validPart) return null
+
+  const result = buildUISurfaces(validPart.messages)
+  const hasInvalidGraph = result.surfaces.some((surface) => getSurfaceGraphError(surface))
+  if (result.error || hasInvalidGraph) {
+    return validPart.fallback ? (
+      <p className="body-sm-regular text-text-tertiary">{validPart.fallback}</p>
+    ) : null
+  }
+
+  if (!result.surfaces.length) return null
+
+  return (
+    <div className="flex w-full flex-col gap-3" data-ui-part-id={validPart.part_id}>
+      {result.surfaces.map((surface) => (
+        <Surface key={surface.surfaceId} surface={surface} />
+      ))}
+    </div>
+  )
+}
+
+export function UIPartList({ parts }: { parts: UIPart[] }) {
+  if (!parts.length) return null
+
+  return (
+    <div className="my-2 flex w-full flex-col gap-3">
+      {parts.map((part) => (
+        <Part key={part.part_id} part={part} />
+      ))}
+    </div>
+  )
+}

web/app/components/base/chat/chat/ui-part/state.ts

@@ -0,0 +1,603 @@
+import type {
+  A2UIBinding,
+  A2UIComponent,
+  A2UIEnvelope,
+  JSONPrimitive,
+  JSONValue,
+  UIPart,
+} from '@/types/a2ui'
+import { A2UI_PROTOCOL_VERSION, DIFY_A2UI_CATALOG_ID } from '@/types/a2ui'
+
+const MAX_MESSAGES_PER_PART = 64
+const MAX_UI_COMPONENTS = 100
+const MAX_JSON_DEPTH = 16
+const MAX_JSON_NODES = 2_000
+const MAX_JSON_ARRAY_INDEX = 1_000
+const MAX_JSON_POINTER_SEGMENTS = 16
+const MAX_STRING_LENGTH = 4_096
+const MAX_PAYLOAD_BYTES = 128 * 1_024
+const MAX_UI_PARTS = 16
+const MAX_UI_PARTS_BYTES = 512 * 1_024
+const MAX_HISTORY_UI_PART_CANDIDATES = 64
+const ID_PATTERN = /^[\dA-Z][\w.:-]{0,127}$/i
+const UNSAFE_PATH_SEGMENTS = new Set(['__proto__', 'constructor', 'prototype'])
+
+type UnknownRecord = Record<string, unknown>
+
+export type UISurface = {
+  surfaceId: string
+  catalogId: typeof DIFY_A2UI_CATALOG_ID
+  components: Map<string, A2UIComponent>
+  dataModel: JSONValue
+}
+
+export type UISurfaceBuildResult = {
+  surfaces: UISurface[]
+  error?: string
+}
+
+function isRecord(value: unknown): value is UnknownRecord {
+  return typeof value === 'object' && value !== null && !Array.isArray(value)
+}
+
+function hasOnlyKeys(
+  value: UnknownRecord,
+  requiredKeys: readonly string[],
+  optionalKeys: readonly string[] = [],
+) {
+  const allowedKeys = new Set([...requiredKeys, ...optionalKeys])
+  return (
+    requiredKeys.every((key) => Object.prototype.hasOwnProperty.call(value, key)) &&
+    Object.keys(value).every((key) => allowedKeys.has(key))
+  )
+}
+
+function isId(value: unknown): value is string {
+  return typeof value === 'string' && ID_PATTERN.test(value)
+}
+
+function isPartId(value: unknown): value is string {
+  return typeof value === 'string' && value.length >= 1 && value.length <= 512
+}
+
+function decodeJSONPointer(value: unknown): string[] | undefined {
+  if (
+    typeof value !== 'string' ||
+    !value.startsWith('/') ||
+    value.length > 256 ||
+    value
+      .split('/')
+      .slice(1)
+      .some((segment) => /~(?![01])/.test(segment))
+  )
+    return undefined
+
+  if (value === '/') return []
+
+  const segments = value
+    .split('/')
+    .slice(1)
+    .map((segment) => segment.replaceAll('~1', '/').replaceAll('~0', '~'))
+  if (segments.length > MAX_JSON_POINTER_SEGMENTS) return undefined
+  return segments.some((segment) => UNSAFE_PATH_SEGMENTS.has(segment)) ? undefined : segments
+}
+
+function isJSONPointer(value: unknown): value is string {
+  return decodeJSONPointer(value) !== undefined
+}
+
+function isBinding(value: unknown): value is A2UIBinding {
+  return isRecord(value) && hasOnlyKeys(value, ['path']) && isJSONPointer(value.path)
+}
+
+function isStringValue(value: unknown): value is string | A2UIBinding {
+  return (typeof value === 'string' && value.length <= MAX_STRING_LENGTH) || isBinding(value)
+}
+
+function isNumberValue(value: unknown): value is number | A2UIBinding {
+  return (typeof value === 'number' && Number.isFinite(value)) || isBinding(value)
+}
+
+function isPrimitive(value: unknown): value is JSONPrimitive {
+  return (
+    value === null ||
+    typeof value === 'string' ||
+    typeof value === 'boolean' ||
+    (typeof value === 'number' && Number.isFinite(value))
+  )
+}
+
+function isPrimitiveValue(value: unknown): value is JSONPrimitive | A2UIBinding {
+  return isPrimitive(value) || isBinding(value)
+}
+
+function isStringArray(value: unknown): value is string[] {
+  return Array.isArray(value) && value.length <= MAX_UI_COMPONENTS && value.every(isId)
+}
+
+function isEnum<T extends string>(value: unknown, allowed: readonly T[]): value is T {
+  return typeof value === 'string' && allowed.includes(value as T)
+}
+
+function isJSONValue(value: unknown): value is JSONValue {
+  let nodeCount = 0
+
+  function visit(current: unknown, depth: number): current is JSONValue {
+    nodeCount += 1
+    if (nodeCount > MAX_JSON_NODES || depth > MAX_JSON_DEPTH) return false
+    if (isPrimitive(current))
+      return typeof current !== 'string' || current.length <= MAX_STRING_LENGTH
+    if (Array.isArray(current)) {
+      for (let index = 0; index < current.length; index++) {
+        if (!visit(current[index], depth + 1)) return false
+      }
+      return true
+    }
+    if (!isRecord(current)) return false
+
+    return Object.entries(current).every(
+      ([key, item]) => !UNSAFE_PATH_SEGMENTS.has(key) && visit(item, depth + 1),
+    )
+  }
+
+  return visit(value, 0)
+}
+
+function isComponent(value: unknown): value is A2UIComponent {
+  if (!isRecord(value) || !isId(value.id) || typeof value.component !== 'string') return false
+
+  switch (value.component) {
+    case 'Card':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'children'], ['title']) &&
+        isStringArray(value.children) &&
+        (value.title === undefined || isStringValue(value.title))
+      )
+    case 'Row':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'children'], ['gap', 'align']) &&
+        isStringArray(value.children) &&
+        (value.gap === undefined || isEnum(value.gap, ['small', 'medium', 'large'] as const)) &&
+        (value.align === undefined || isEnum(value.align, ['start', 'center', 'end'] as const))
+      )
+    case 'Column':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'children'], ['gap']) &&
+        isStringArray(value.children) &&
+        (value.gap === undefined || isEnum(value.gap, ['small', 'medium', 'large'] as const))
+      )
+    case 'Text':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'text'], ['variant']) &&
+        isStringValue(value.text) &&
+        (value.variant === undefined || isEnum(value.variant, ['body', 'caption'] as const))
+      )
+    case 'Icon':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'name']) &&
+        isEnum(value.name, [
+          'clock',
+          'cloud',
+          'sun',
+          'rain',
+          'snow',
+          'wind',
+          'thermometer',
+          'calendar',
+          'location',
+        ] as const)
+      )
+    case 'Divider':
+      return hasOnlyKeys(value, ['id', 'component'])
+    case 'Badge':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'text'], ['tone']) &&
+        isStringValue(value.text) &&
+        (value.tone === undefined ||
+          isEnum(value.tone, ['neutral', 'info', 'success', 'warning', 'critical'] as const))
+      )
+    case 'Metric':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'label', 'value'], ['unit']) &&
+        isStringValue(value.label) &&
+        isPrimitiveValue(value.value) &&
+        (value.unit === undefined || isStringValue(value.unit))
+      )
+    case 'DateTime':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'value'], ['format']) &&
+        isStringValue(value.value) &&
+        (value.format === undefined || isEnum(value.format, ['date', 'time', 'datetime'] as const))
+      )
+    case 'Progress':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'value', 'label'], ['max']) &&
+        isNumberValue(value.value) &&
+        (value.max === undefined || isNumberValue(value.max)) &&
+        isStringValue(value.label)
+      )
+    case 'KeyValue':
+      return (
+        hasOnlyKeys(value, ['id', 'component', 'label', 'value']) &&
+        isStringValue(value.label) &&
+        isPrimitiveValue(value.value)
+      )
+    default:
+      return false
+  }
+}
+
+function isCreateSurfaceEnvelope(value: UnknownRecord): value is A2UIEnvelope {
+  if (!hasOnlyKeys(value, ['version', 'createSurface'])) return false
+  if (!isRecord(value.createSurface)) return false
+
+  return (
+    hasOnlyKeys(value.createSurface, ['surfaceId', 'catalogId']) &&
+    isId(value.createSurface.surfaceId) &&
+    value.createSurface.catalogId === DIFY_A2UI_CATALOG_ID
+  )
+}
+
+function isUpdateComponentsEnvelope(value: UnknownRecord): value is A2UIEnvelope {
+  if (!hasOnlyKeys(value, ['version', 'updateComponents'])) return false
+  if (!isRecord(value.updateComponents)) return false
+
+  const components = value.updateComponents.components
+  if (
+    !hasOnlyKeys(value.updateComponents, ['surfaceId', 'components']) ||
+    !isId(value.updateComponents.surfaceId) ||
+    !Array.isArray(components) ||
+    components.length === 0 ||
+    components.length > MAX_UI_COMPONENTS
+  )
+    return false
+
+  const componentIds = new Set<string>()
+  return components.every((component) => {
+    if (!isComponent(component) || componentIds.has(component.id)) return false
+    componentIds.add(component.id)
+    return true
+  })
+}
+
+function isUpdateDataModelEnvelope(value: UnknownRecord): value is A2UIEnvelope {
+  if (!hasOnlyKeys(value, ['version', 'updateDataModel'])) return false
+  if (!isRecord(value.updateDataModel)) return false
+
+  return (
+    hasOnlyKeys(value.updateDataModel, ['surfaceId', 'value'], ['path']) &&
+    isId(value.updateDataModel.surfaceId) &&
+    (value.updateDataModel.path === undefined || isJSONPointer(value.updateDataModel.path)) &&
+    isJSONValue(value.updateDataModel.value)
+  )
+}
+
+function isDeleteSurfaceEnvelope(value: UnknownRecord): value is A2UIEnvelope {
+  if (!hasOnlyKeys(value, ['version', 'deleteSurface'])) return false
+  if (!isRecord(value.deleteSurface)) return false
+
+  return hasOnlyKeys(value.deleteSurface, ['surfaceId']) && isId(value.deleteSurface.surfaceId)
+}
+
+function isEnvelope(value: unknown): value is A2UIEnvelope {
+  if (!isRecord(value) || value.version !== A2UI_PROTOCOL_VERSION) return false
+
+  return (
+    isCreateSurfaceEnvelope(value) ||
+    isUpdateComponentsEnvelope(value) ||
+    isUpdateDataModelEnvelope(value) ||
+    isDeleteSurfaceEnvelope(value)
+  )
+}
+
+function getEnvelopeSurfaceId(message: A2UIEnvelope) {
+  if ('createSurface' in message) return message.createSurface.surfaceId
+  if ('updateComponents' in message) return message.updateComponents.surfaceId
+  if ('updateDataModel' in message) return message.updateDataModel.surfaceId
+  return message.deleteSurface.surfaceId
+}
+
+function hasValidMessageLifecycle(messages: A2UIEnvelope[]) {
+  const firstMessage = messages[0]
+  if (!firstMessage || !('createSurface' in firstMessage)) return false
+
+  const surfaceId = firstMessage.createSurface.surfaceId
+  let deleteSeen = false
+  for (const [index, message] of messages.entries()) {
+    if (getEnvelopeSurfaceId(message) !== surfaceId) return false
+    if (index > 0 && 'createSurface' in message) return false
+    if ('deleteSurface' in message) {
+      if (deleteSeen || index !== messages.length - 1) return false
+      deleteSeen = true
+    }
+  }
+  return true
+}
+
+function getUIPartPayloadByteLength(value: UnknownRecord) {
+  try {
+    return new TextEncoder().encode(
+      JSON.stringify({
+        protocol: value.protocol,
+        protocol_version: value.protocol_version,
+        messages: value.messages,
+        fallback: value.fallback,
+      }),
+    ).byteLength
+  } catch {
+    return Number.POSITIVE_INFINITY
+  }
+}
+
+export function parseUIPartEnvelope(value: unknown): UIPart | undefined {
+  if (!isRecord(value)) return undefined
+  if (
+    !hasOnlyKeys(
+      value,
+      ['part_id', 'sequence', 'protocol', 'protocol_version', 'messages'],
+      ['fallback'],
+    )
+  )
+    return undefined
+
+  if (
+    !isPartId(value.part_id) ||
+    !Number.isSafeInteger(value.sequence) ||
+    (value.sequence as number) < 1 ||
+    value.protocol !== 'a2ui' ||
+    value.protocol_version !== A2UI_PROTOCOL_VERSION ||
+    !Array.isArray(value.messages) ||
+    value.messages.length === 0 ||
+    value.messages.length > MAX_MESSAGES_PER_PART ||
+    !value.messages.every(isEnvelope) ||
+    (value.fallback !== undefined &&
+      value.fallback !== null &&
+      (typeof value.fallback !== 'string' || value.fallback.length > MAX_STRING_LENGTH)) ||
+    getUIPartPayloadByteLength(value) > MAX_PAYLOAD_BYTES
+  )
+    return undefined
+
+  return value as UIPart
+}
+
+export function parseUIPart(value: unknown): UIPart | undefined {
+  const part = parseUIPartEnvelope(value)
+  if (!part) return undefined
+  if (!hasValidMessageLifecycle(part.messages)) return undefined
+
+  return buildUISurfaces(part.messages).error ? undefined : part
+}
+
+type UIPartUpsertResult = {
+  parts: UIPart[]
+  resourceLimitExceeded: boolean
+}
+
+function getUIPartsByteLength(parts: UIPart[]) {
+  try {
+    return new TextEncoder().encode(JSON.stringify(parts)).byteLength
+  } catch {
+    return Number.POSITIVE_INFINITY
+  }
+}
+
+function upsertParsedUIPart(current: UIPart[], part: UIPart): UIPartUpsertResult {
+  if (current.length > MAX_UI_PARTS || getUIPartsByteLength(current) > MAX_UI_PARTS_BYTES)
+    return { parts: current, resourceLimitExceeded: true }
+
+  const existingIndex = current.findIndex((item) => item.part_id === part.part_id)
+  if (existingIndex === -1 && current.length >= MAX_UI_PARTS)
+    return { parts: current, resourceLimitExceeded: true }
+  if (existingIndex > -1 && current[existingIndex]!.sequence >= part.sequence)
+    return { parts: current, resourceLimitExceeded: false }
+
+  const next = [...current]
+  if (existingIndex === -1) next.push(part)
+  else next[existingIndex] = part
+
+  return getUIPartsByteLength(next) > MAX_UI_PARTS_BYTES
+    ? { parts: current, resourceLimitExceeded: true }
+    : { parts: next, resourceLimitExceeded: false }
+}
+
+export function upsertUIPart(current: UIPart[], incoming: unknown): UIPart[] {
+  const part = parseUIPart(incoming)
+  if (!part) return current
+
+  return upsertParsedUIPart(current, part).parts
+}
+
+function normalizeUIPartsWithStatus(value: unknown): UIPartUpsertResult {
+  if (!Array.isArray(value)) return { parts: [], resourceLimitExceeded: false }
+  if (value.length > MAX_HISTORY_UI_PART_CANDIDATES)
+    return { parts: [], resourceLimitExceeded: true }
+
+  let parts: UIPart[] = []
+  for (const item of value) {
+    const part = parseUIPart(item)
+    if (!part) continue
+
+    const result = upsertParsedUIPart(parts, part)
+    if (result.resourceLimitExceeded) return { parts: [], resourceLimitExceeded: true }
+    parts = result.parts
+  }
+  return { parts, resourceLimitExceeded: false }
+}
+
+export function normalizeUIParts(value: unknown): UIPart[] {
+  return normalizeUIPartsWithStatus(value).parts
+}
+
+export function reconcileHistoryUIParts(
+  current: UIPart[] | undefined,
+  historyValue: unknown,
+): UIPart[] | undefined {
+  if (historyValue === undefined) return current
+
+  const normalized = normalizeUIPartsWithStatus(historyValue)
+  return normalized.resourceLimitExceeded ? current : normalized.parts
+}
+
+export function buildUISurfaces(messages: A2UIEnvelope[]): UISurfaceBuildResult {
+  if (!messages.every(isEnvelope)) return { surfaces: [], error: 'invalid_messages' }
+
+  const surfaces = new Map<string, UISurface>()
+
+  for (const message of messages) {
+    if ('createSurface' in message) {
+      const { surfaceId, catalogId } = message.createSurface
+      surfaces.set(surfaceId, {
+        surfaceId,
+        catalogId,
+        components: new Map(),
+        dataModel: {},
+      })
+      continue
+    }
+
+    if ('updateComponents' in message) {
+      const surface = surfaces.get(message.updateComponents.surfaceId)
+      if (!surface) continue
+      for (const component of message.updateComponents.components) {
+        surface.components.set(component.id, component)
+        if (surface.components.size > MAX_UI_COMPONENTS)
+          return { surfaces: [], error: 'too_many_components' }
+      }
+      continue
+    }
+
+    if ('updateDataModel' in message) {
+      const surface = surfaces.get(message.updateDataModel.surfaceId)
+      if (surface) {
+        const updatedModel = updateJSONPointer(
+          surface.dataModel,
+          message.updateDataModel.path ?? '/',
+          message.updateDataModel.value,
+        )
+        if (updatedModel === undefined || !isJSONValue(updatedModel))
+          return { surfaces: [], error: 'invalid_data_model_update' }
+        surface.dataModel = updatedModel
+      }
+      continue
+    }
+
+    surfaces.delete(message.deleteSurface.surfaceId)
+  }
+
+  const builtSurfaces = [...surfaces.values()]
+  for (const surface of builtSurfaces) {
+    const graphError = getSurfaceGraphError(surface)
+    if (graphError) return { surfaces: [], error: graphError }
+  }
+
+  return { surfaces: builtSurfaces }
+}
+
+export function getSurfaceGraphError(surface: UISurface): string | undefined {
+  if (surface.components.size > MAX_UI_COMPONENTS) return 'too_many_components'
+  if (!surface.components.has('root')) return 'missing_root'
+
+  const parents = new Map<string, string>()
+  for (const component of surface.components.values()) {
+    if (
+      component.component !== 'Card' &&
+      component.component !== 'Row' &&
+      component.component !== 'Column'
+    )
+      continue
+
+    const ownChildren = new Set<string>()
+    for (const childId of component.children) {
+      if (ownChildren.has(childId)) return 'duplicate_child'
+      ownChildren.add(childId)
+      if (childId === 'root') return 'root_as_child'
+      if (!surface.components.has(childId)) return 'missing_component'
+      if (parents.has(childId)) return 'multiple_parents'
+      parents.set(childId, component.id)
+    }
+  }
+
+  const visited = new Set<string>()
+  const pending = ['root']
+  while (pending.length) {
+    const componentId = pending.pop()!
+    if (visited.has(componentId)) return 'component_cycle'
+    const component = surface.components.get(componentId)
+    if (!component) return 'missing_component'
+
+    visited.add(componentId)
+    if (visited.size > MAX_UI_COMPONENTS) return 'render_budget_exceeded'
+    if (
+      component.component === 'Card' ||
+      component.component === 'Row' ||
+      component.component === 'Column'
+    )
+      pending.push(...component.children)
+  }
+
+  return visited.size === surface.components.size ? undefined : 'unreachable_component'
+}
+
+export function resolveJSONPointer(model: JSONValue, pointer: string): JSONValue | undefined {
+  const segments = decodeJSONPointer(pointer)
+  if (!segments) return undefined
+  if (!segments.length) return model
+
+  let current: JSONValue | undefined = model
+
+  for (const segment of segments) {
+    if (Array.isArray(current)) {
+      if (!/^(?:0|[1-9]\d*)$/.test(segment)) return undefined
+      current = current[Number(segment)]
+      continue
+    }
+    if (isRecord(current)) {
+      const next = current[segment]
+      current = isJSONValue(next) ? next : undefined
+      continue
+    }
+    return undefined
+  }
+
+  return current
+}
+
+function updateJSONPointer(
+  model: JSONValue,
+  pointer: string,
+  value: JSONValue,
+): JSONValue | undefined {
+  const segments = decodeJSONPointer(pointer)
+  if (!segments) return undefined
+  if (!segments.length) return value
+  const pointerSegments = segments
+
+  function update(current: JSONValue | undefined, segmentIndex: number): JSONValue | undefined {
+    if (segmentIndex === pointerSegments.length) return value
+
+    const segment = pointerSegments[segmentIndex]!
+    const isArraySegment = /^(?:0|[1-9]\d*)$/.test(segment)
+    if (Array.isArray(current) || (!isRecord(current) && isArraySegment)) {
+      if (!isArraySegment) return undefined
+      const index = Number(segment)
+      if (index > MAX_JSON_ARRAY_INDEX) return undefined
+
+      const next = Array.isArray(current) ? [...current] : []
+      if (index > next.length) return undefined
+      const updatedChild = update(next[index], segmentIndex + 1)
+      if (updatedChild === undefined) return undefined
+      next[index] = updatedChild
+      return next
+    }
+
+    const next: Record<string, JSONValue> = isRecord(current)
+      ? { ...(current as Record<string, JSONValue>) }
+      : {}
+    const updatedChild = update(next[segment], segmentIndex + 1)
+    if (updatedChild === undefined) return undefined
+    next[segment] = updatedChild
+    return next
+  }
+
+  return update(model, 0)
+}

web/app/components/base/chat/embedded-chatbot/__tests__/hooks.spec.tsx

@@ -1,6 +1,7 @@
 import type { ReactNode } from 'react'
 import type { ChatConfig } from '../../types'
 import type { AppConversationData, AppData, AppMeta, ConversationItem } from '@/models/share'
+import type { UIPart } from '@/types/a2ui'
 import { ToastHost } from '@langgenius/dify-ui/toast'
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import { act, renderHook, waitFor } from '@testing-library/react'
@@ -153,6 +154,29 @@ const createConversationData = (
   ...overrides,
 })
 
+const createHistoricalUIPart = (): UIPart => ({
+  part_id: 'call/abc=:weather',
+  sequence: 1,
+  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: 'Card', children: [], title: 'Weather' }],
+      },
+    },
+  ],
+})
+
 // Scenario: useEmbeddedChatbot integrates share queries for conversations and chat list.
 describe('useEmbeddedChatbot', () => {
   beforeEach(() => {
@@ -293,6 +317,35 @@ describe('useEmbeddedChatbot', () => {
         ((assistantMsg as Record<string, unknown>)?.feedback as Record<string, unknown>)?.rating,
       ).toBe('like')
     })
+
+    it('should restore valid UI parts from message history metadata', async () => {
+      const uiPart = createHistoricalUIPart()
+      mockGetProcessedSystemVariablesFromUrlParams.mockResolvedValue({
+        conversation_id: 'conversation-1',
+      })
+      mockFetchChatList.mockResolvedValue({
+        data: [
+          {
+            id: 'msg-with-ui',
+            query: 'What is the weather?',
+            answer: '',
+            message_files: [],
+            agent_thoughts: null,
+            metadata: {
+              ui_parts: [uiPart],
+            },
+          },
+        ],
+      })
+
+      const { result } = await renderWithClient(() => useEmbeddedChatbot(AppSourceType.webApp))
+
+      await waitFor(() => {
+        expect(result.current.appPrevChatList).toHaveLength(1)
+      })
+      const answerNode = result.current.appPrevChatList[0]?.children?.[0]
+      expect(answerNode?.ui_parts).toEqual([uiPart])
+    })
   })
 
   // Scenario: completion invalidates share caches and merges generated names.

web/app/components/base/chat/embedded-chatbot/hooks.tsx

@@ -23,6 +23,7 @@ import {
 import { useGetTryAppInfo, useGetTryAppParams } from '@/service/use-try-app'
 import { TransferMethod } from '@/types/app'
 import { getProcessedFilesFromResponse } from '../../file-uploader/utils'
+import { normalizeUIParts } from '../chat/ui-part/state'
 import {
   buildChatItemTree,
   getProcessedInputsFromUrlParams,
@@ -58,6 +59,7 @@ function getFormattedChatList(messages: any[]) {
       citation: item.retriever_resources,
       reasoningContent: item.metadata?.reasoning,
       reasoningFinished: true,
+      ui_parts: normalizeUIParts(item.ui_parts ?? item.metadata?.ui_parts),
       message_files: getProcessedFilesFromResponse(
         answerFiles.map((item: any) => ({ ...item, related_id: item.id })),
       ),

web/app/components/base/markdown/__tests__/streamdown-wrapper.spec.tsx

@@ -1,5 +1,6 @@
 import type { PropsWithChildren, ReactNode } from 'react'
 import { render, screen } from '@testing-library/react'
+import { preprocessThinkTag } from '../markdown-utils'
 import StreamdownWrapper from '../streamdown-wrapper'
 
 const TILDE_RANGE_RE = /0\.3~8mm/
@@ -141,6 +142,27 @@ describe('StreamdownWrapper', () => {
     })
   })
 
+  describe('Think block rendering', () => {
+    it.each([
+      ['immediately', ''],
+      ['after one newline', '\n'],
+    ])(
+      'should render markdown following a think block outside the details element %s',
+      (_, separator) => {
+        const content = preprocessThinkTag(`<think>Reasoning</think>${separator}**Final answer**`)
+
+        render(<StreamdownWrapper latexContent={content} />)
+
+        const details = document.querySelector('details')
+        expect(details).not.toBeNull()
+        expect(details).toHaveTextContent('Reasoning')
+        expect(details).not.toHaveTextContent('Final answer')
+        expect(screen.getByText('Final answer')).toBeInTheDocument()
+        expect(document.querySelector('[data-streamdown="strong"]')).not.toBeNull()
+      },
+    )
+  })
+
   describe('Plugin Info behavior', () => {
     it('should render PluginImg and PluginParagraph when pluginInfo is provided', () => {
       // Arrange

web/app/components/base/markdown/markdown-utils.ts

@@ -36,8 +36,7 @@ export const preprocessThinkTag = (content: string) => {
   const thinkCloseTagRegex = /(\s*<\/think>)+/g
   return flow([
     (str: string) => str.replace(thinkOpenTagRegex, '<details data-think=true>\n'),
-    (str: string) => str.replace(thinkCloseTagRegex, '\n[ENDTHINKFLAG]</details>'),
-    (str: string) => str.replace(/(<\/details>)(?![^\S\r\n]*[\r\n])(?![^\S\r\n]*$)/g, '$1\n'),
+    (str: string) => str.replace(thinkCloseTagRegex, '\n[ENDTHINKFLAG]</details>\n\n'),
   ])(content)
 }
 

web/app/components/workflow/__tests__/workflow-edge-events.spec.tsx

@@ -165,6 +165,7 @@ vi.mock('@/next/navigation', () => ({
   useRouter: () => ({
     push: vi.fn(),
   }),
+  useSearchParams: () => new URLSearchParams(),
 }))
 
 vi.mock('@/context/event-emitter', () => ({
@@ -445,6 +446,10 @@ vi.mock('../hooks/use-workflow-search', () => ({
   useWorkflowSearch: workflowHookMocks.useWorkflowSearch,
 }))
 
+vi.mock('../hooks/use-locate-node', () => ({
+  useLocateNode: vi.fn(),
+}))
+
 vi.mock('../nodes/_base/components/variable/use-match-schema-type', () => ({
   default: () => ({
     schemaTypeDefinitions: undefined,

web/app/components/workflow/hooks/__tests__/use-locate-node.spec.tsx

@@ -0,0 +1,176 @@
+import type { CommonNodeType, Node } from '../../types'
+import { renderHook } from '@testing-library/react'
+import { useLocateNode } from '../use-locate-node'
+
+const mockHandleNodeSelect = vi.hoisted(() => vi.fn())
+const mockScrollToWorkflowNode = vi.hoisted(() => vi.fn())
+const mockSearchParams = vi.hoisted(() => new URLSearchParams())
+const mockToastSuccess = vi.hoisted(() => vi.fn())
+const mockToastError = vi.hoisted(() => vi.fn())
+
+vi.mock('@/next/navigation', () => ({
+  useSearchParams: () => mockSearchParams,
+}))
+
+vi.mock('@langgenius/dify-ui/toast', () => ({
+  toast: Object.assign(vi.fn(), {
+    success: mockToastSuccess,
+    error: mockToastError,
+    warning: vi.fn(),
+    info: vi.fn(),
+    dismiss: vi.fn(),
+    update: vi.fn(),
+    promise: vi.fn(),
+  }),
+}))
+
+vi.mock('react-i18next', () => ({
+  useTranslation: () => ({
+    t: (
+      selector: (source: Record<string, string>) => string,
+      options?: Record<string, unknown>,
+    ) => {
+      const key = selector(
+        new Proxy({}, { get: (_, property: string) => property }) as Record<string, string>,
+      )
+      if (options?.nodeId) return `${key}:${options.nodeId}`
+      if (options?.title) return `${key}:${options.title}`
+      return key
+    },
+  }),
+}))
+
+vi.mock('../use-nodes-interactions', () => ({
+  useNodesInteractions: () => ({
+    handleNodeSelect: mockHandleNodeSelect,
+  }),
+}))
+
+vi.mock('../../utils/node-navigation', () => ({
+  scrollToWorkflowNode: (nodeId: string) => mockScrollToWorkflowNode(nodeId),
+}))
+
+const createNode = (overrides: Partial<Node> = {}): Node => ({
+  id: 'node-1',
+  type: 'custom',
+  position: { x: 0, y: 0 },
+  data: {
+    type: 'llm',
+    title: 'Writer',
+    desc: 'Draft content',
+  } as CommonNodeType,
+  ...overrides,
+})
+
+describe('useLocateNode', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    vi.useFakeTimers()
+    mockSearchParams.delete('node_id')
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  it('does nothing when node_id param is absent', () => {
+    const nodes = [createNode({ id: 'n1' })]
+    renderHook(() => useLocateNode(nodes))
+
+    expect(mockHandleNodeSelect).not.toHaveBeenCalled()
+    expect(mockToastSuccess).not.toHaveBeenCalled()
+    expect(mockToastError).not.toHaveBeenCalled()
+  })
+
+  it('does nothing when nodes are empty', () => {
+    mockSearchParams.set('node_id', 'n1')
+    renderHook(() => useLocateNode([]))
+
+    expect(mockHandleNodeSelect).not.toHaveBeenCalled()
+    expect(mockToastError).not.toHaveBeenCalled()
+  })
+
+  it('selects, scrolls to, and shows success toast when node is found', () => {
+    mockSearchParams.set('node_id', 'target-node')
+    const nodes = [createNode({ id: 'target-node', data: { title: 'My Node' } as CommonNodeType })]
+
+    renderHook(() => useLocateNode(nodes))
+
+    expect(mockHandleNodeSelect).toHaveBeenCalledWith('target-node')
+    expect(mockToastSuccess).toHaveBeenCalledWith('panel.locateNodeSuccess:My Node')
+
+    // Scroll happens after a 200ms delay
+    expect(mockScrollToWorkflowNode).not.toHaveBeenCalled()
+    vi.advanceTimersByTime(200)
+    expect(mockScrollToWorkflowNode).toHaveBeenCalledWith('target-node')
+  })
+
+  it('shows error toast after debounce when node is not found', () => {
+    mockSearchParams.set('node_id', 'missing-node')
+    const nodes = [createNode({ id: 'other-node' })]
+
+    renderHook(() => useLocateNode(nodes))
+
+    // Not immediately reported
+    expect(mockToastError).not.toHaveBeenCalled()
+
+    // After 500ms debounce, reports not found
+    vi.advanceTimersByTime(500)
+    expect(mockToastError).toHaveBeenCalledWith('panel.locateNodeNotFound:missing-node')
+  })
+
+  it('does not report not-found prematurely when nodes are still loading', () => {
+    mockSearchParams.set('node_id', 'target-node')
+    const initialNodes = [createNode({ id: 'other-node' })]
+
+    const { rerender } = renderHook(({ nodes }) => useLocateNode(nodes), {
+      initialProps: { nodes: initialNodes },
+    })
+
+    // Before debounce fires, more nodes arrive including the target
+    vi.advanceTimersByTime(300)
+    rerender({
+      nodes: [
+        createNode({ id: 'other-node' }),
+        createNode({ id: 'target-node', data: { title: 'Found!' } as CommonNodeType }),
+      ],
+    })
+
+    // Should have located the node, not reported error
+    expect(mockToastError).not.toHaveBeenCalled()
+    expect(mockHandleNodeSelect).toHaveBeenCalledWith('target-node')
+    expect(mockToastSuccess).toHaveBeenCalledWith('panel.locateNodeSuccess:Found!')
+  })
+
+  it('locates only once even if nodes update afterwards', () => {
+    mockSearchParams.set('node_id', 'target-node')
+    const nodes1 = [createNode({ id: 'target-node' })]
+
+    const { rerender } = renderHook(({ nodes }) => useLocateNode(nodes), {
+      initialProps: { nodes: nodes1 },
+    })
+
+    expect(mockHandleNodeSelect).toHaveBeenCalledTimes(1)
+
+    // Simulate a nodes update after locate has already happened
+    rerender({
+      nodes: [createNode({ id: 'target-node' }), createNode({ id: 'extra' })],
+    })
+
+    expect(mockHandleNodeSelect).toHaveBeenCalledTimes(1)
+    expect(mockToastSuccess).toHaveBeenCalledTimes(1)
+  })
+
+  it('clears scroll timeout on unmount', () => {
+    mockSearchParams.set('node_id', 'target-node')
+    const nodes = [createNode({ id: 'target-node' })]
+
+    const { unmount } = renderHook(() => useLocateNode(nodes))
+
+    // Unmount before the 200ms scroll timer fires
+    unmount()
+
+    vi.advanceTimersByTime(500)
+    expect(mockScrollToWorkflowNode).not.toHaveBeenCalled()
+  })
+})

web/app/components/workflow/hooks/__tests__/use-workflow-search.spec.tsx

@@ -94,13 +94,68 @@ describe('useWorkflowSearch', () => {
 
     const toolResults = findWorkflowNodes('search')
     expect(toolResults.map((item) => item.id)).toEqual(['tool-1'])
-    expect(toolResults[0]?.description).toBe('Search the web')
+    // description now includes nodeId suffix: "desc · nodeId"
+    expect(toolResults[0]?.description).toBe('Search the web · tool-1')
 
     unmount()
 
     expect(findWorkflowNodes('gpt')).toEqual([])
   })
 
+  it('matches by node_id with highest priority scoring', async () => {
+    runtimeNodes.push(
+      createNode({
+        id: '1721234567890',
+        data: {
+          type: BlockEnum.LLM,
+          title: 'Writer',
+          desc: 'Draft content',
+        } as CommonNodeType,
+      }),
+      createNode({
+        id: 'other-node',
+        data: {
+          type: BlockEnum.LLM,
+          title: '1721234567890', // title also matches the searchTerm
+          desc: '',
+        } as CommonNodeType,
+      }),
+    )
+
+    const { unmount } = renderHook(() => useWorkflowSearch())
+
+    // Exact nodeId match (120pts) should rank higher than title exact match (100pts)
+    const results = findWorkflowNodes('1721234567890')
+    expect(results.map((item) => item.id)).toEqual(['1721234567890', 'other-node'])
+
+    unmount()
+  })
+
+  it('matches by partial node_id', async () => {
+    runtimeNodes.push(
+      createNode({
+        id: '1721234567890',
+        data: {
+          type: BlockEnum.LLM,
+          title: 'Writer',
+          desc: 'Draft content',
+        } as CommonNodeType,
+      }),
+    )
+
+    const { unmount } = renderHook(() => useWorkflowSearch())
+
+    // Prefix match
+    const prefixResults = findWorkflowNodes('172123')
+    expect(prefixResults.map((item) => item.id)).toEqual(['1721234567890'])
+
+    // Partial match
+    const partialResults = findWorkflowNodes('456')
+    expect(partialResults.map((item) => item.id)).toEqual(['1721234567890'])
+
+    unmount()
+  })
+
   it('binds the node selection listener to handleNodeSelect', () => {
     const { unmount } = renderHook(() => useWorkflowSearch())
 

web/app/components/workflow/hooks/use-locate-node.ts

@@ -0,0 +1,80 @@
+'use client'
+
+import type { Node } from '../types'
+import { toast } from '@langgenius/dify-ui/toast'
+import { useEffect, useRef } from 'react'
+import { useTranslation } from 'react-i18next'
+import { useSearchParams } from '@/next/navigation'
+import { scrollToWorkflowNode } from '../utils/node-navigation'
+import { useNodesInteractions } from './use-nodes-interactions'
+
+/**
+ * Hook to locate a node by ID from URL query parameter `node_id`.
+ *
+ * Usage scenario: operators find a failing node ID from server logs,
+ * construct a URL like `/workflow?node_id=xxx`, and the editor will
+ * automatically select and scroll to that node on load.
+ *
+ * The hook reads `node_id` from the URL search params, waits for nodes
+ * to be available, then selects and scrolls to the target node.
+ * A toast message is shown to indicate success or failure.
+ */
+export const useLocateNode = (nodes: Node[]) => {
+  const { t } = useTranslation('workflow')
+  const searchParams = useSearchParams()
+  const nodeIdFromUrl = searchParams.get('node_id')
+  const { handleNodeSelect } = useNodesInteractions()
+  const hasLocateRef = useRef(false)
+
+  useEffect(() => {
+    if (!nodeIdFromUrl || hasLocateRef.current) return
+
+    // Wait for nodes to be loaded
+    if (!nodes.length) return
+
+    const targetNode = nodes.find((n) => n.id === nodeIdFromUrl)
+
+    if (!targetNode) {
+      // Don't mark as located yet — nodes may still be loading asynchronously.
+      // Retry on the next nodes update before reporting "not found".
+      return
+    }
+
+    // Select the node (opens its panel) and scroll to it
+    handleNodeSelect(nodeIdFromUrl)
+
+    // Delay scroll to ensure node selection state has been applied
+    const scrollTimer = setTimeout(() => {
+      scrollToWorkflowNode(nodeIdFromUrl)
+    }, 200)
+
+    toast.success(
+      t(($) => $['panel.locateNodeSuccess'], {
+        title: targetNode.data?.title || nodeIdFromUrl,
+      }),
+    )
+    hasLocateRef.current = true
+
+    return () => clearTimeout(scrollTimer)
+  }, [nodeIdFromUrl, nodes, handleNodeSelect, t])
+
+  // Report "not found" after nodes have settled (no longer changing)
+  useEffect(() => {
+    if (!nodeIdFromUrl || hasLocateRef.current || !nodes.length) return
+
+    const targetNode = nodes.find((n) => n.id === nodeIdFromUrl)
+    if (targetNode) return
+
+    // Debounce: wait to see if nodes continue loading before reporting not found
+    const notFoundTimer = setTimeout(() => {
+      if (hasLocateRef.current) return
+      const stillMissing = !nodes.find((n) => n.id === nodeIdFromUrl)
+      if (stillMissing) {
+        toast.error(t(($) => $['panel.locateNodeNotFound'], { nodeId: nodeIdFromUrl }))
+        hasLocateRef.current = true
+      }
+    }, 500)
+
+    return () => clearTimeout(notFoundTimer)
+  }, [nodeIdFromUrl, nodes, t])
+}

web/app/components/workflow/hooks/use-workflow-search.tsx

@@ -80,6 +80,7 @@ export const useWorkflowSearch = () => {
 
       return {
         id: node.id,
+        nodeId: node.id,
         title: nodeData?.title || nodeData?.type || 'Untitled',
         type: nodeData?.type || '',
         desc: nodeData?.desc || '',
@@ -95,6 +96,7 @@ export const useWorkflowSearch = () => {
   const calculateScore = useCallback(
     (
       node: {
+        nodeId: string
         title: string
         type: string
         desc: string
@@ -107,12 +109,18 @@ export const useWorkflowSearch = () => {
       const titleMatch = node.title.toLowerCase()
       const typeMatch = node.type.toLowerCase()
       const descMatch = node.desc?.toLowerCase() || ''
+      const nodeIdMatch = node.nodeId?.toLowerCase() || ''
       const modelProviderMatch = node.modelInfo?.provider?.toLowerCase() || ''
       const modelNameMatch = node.modelInfo?.name?.toLowerCase() || ''
       const modelModeMatch = node.modelInfo?.mode?.toLowerCase() || ''
 
       let score = 0
 
+      // Node ID matching (exact > partial — useful for locating nodes from server logs)
+      if (nodeIdMatch === searchTerm) score += 120
+      else if (nodeIdMatch.startsWith(searchTerm)) score += 90
+      else if (nodeIdMatch.includes(searchTerm)) score += 40
+
       // Title matching (exact prefix > partial match)
       if (titleMatch.startsWith(searchTerm)) score += 100
       else if (titleMatch.includes(searchTerm)) score += 50
@@ -149,7 +157,7 @@ export const useWorkflowSearch = () => {
             ? {
                 id: node.id,
                 title: node.title,
-                description: node.desc || node.type,
+                description: [node.desc || node.type, node.nodeId].filter(Boolean).join(' · '),
                 type: 'workflow-node' as const,
                 path: `#${node.id}`,
                 icon: (