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

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/service/base.spec.ts

@@ -316,6 +316,53 @@ describe('handleStream', () => {
       expect(onData).not.toHaveBeenCalled()
     })
 
+    it('should dispatch ui_part events without treating them as text', async () => {
+      const onData = vi.fn()
+      const onUIPart = vi.fn()
+      const uiPartEvent = {
+        event: 'ui_part',
+        id: 'message-1',
+        part: {
+          part_id: 'weather',
+          sequence: 1,
+          protocol: 'a2ui',
+          protocol_version: 'v0.9.1',
+          messages: [],
+        },
+      }
+      const mockReader = {
+        read: vi
+          .fn()
+          .mockResolvedValueOnce({
+            done: false,
+            value: new TextEncoder().encode(`data: ${JSON.stringify(uiPartEvent)}\n`),
+          })
+          .mockResolvedValueOnce({
+            done: true,
+            value: undefined,
+          }),
+      }
+      const mockResponse = {
+        ok: true,
+        body: {
+          getReader: () => mockReader,
+        },
+      } as unknown as Response
+      const callbacksBeforeUIPart = Array.from({ length: 32 }, () => undefined)
+
+      ;(handleStream as (...args: unknown[]) => void)(
+        mockResponse,
+        onData,
+        ...callbacksBeforeUIPart,
+        onUIPart,
+      )
+
+      await waitFor(() => {
+        expect(onUIPart).toHaveBeenCalledWith(uiPartEvent)
+      })
+      expect(onData).not.toHaveBeenCalled()
+    })
+
     it('should complete with error when the stream reader rejects', async () => {
       const onData = vi.fn()
       const onCompleted = vi.fn()

web/service/base.ts

@@ -1,5 +1,6 @@
 import type { FetchOptionType, ResponseError } from './fetch'
 import type { MessageEnd, MessageReplace, ThoughtItem } from '@/app/components/base/chat/chat/type'
+import type { UIPartStreamEvent } from '@/types/a2ui'
 import type { VisionFile } from '@/types/app'
 import type {
   DataSourceNodeCompletedResponse,
@@ -92,6 +93,7 @@ type IOnLoopStarted = (workflowStarted: LoopStartedResponse) => void
 type IOnLoopNext = (workflowStarted: LoopNextResponse) => void
 type IOnLoopFinished = (workflowFinished: LoopFinishedResponse) => void
 type IOnAgentLog = (agentLog: AgentLogResponse) => void
+type IOnUIPart = (uiPart: UIPartStreamEvent) => void
 
 type IOHumanInputRequired = (humanInputRequired: HumanInputRequiredResponse) => void
 type IOnHumanInputFormFilled = (humanInputFormFilled: HumanInputFormFilledResponse) => void
@@ -144,6 +146,7 @@ export type IOtherOptions = {
   onLoopNext?: IOnLoopNext
   onLoopFinish?: IOnLoopFinished
   onAgentLog?: IOnAgentLog
+  onUIPart?: IOnUIPart
   onHumanInputRequired?: IOHumanInputRequired
   onHumanInputFormFilled?: IOnHumanInputFormFilled
   onHumanInputFormTimeout?: IOnHumanInputFormTimeout
@@ -275,6 +278,7 @@ export const handleStream = (
   onDataSourceNodeError?: IOnDataSourceNodeError,
   onReasoning?: IOnReasoning,
   onUnhandledEvent?: IOnUnhandledEvent,
+  onUIPart?: IOnUIPart,
 ) => {
   if (!response.ok) throw new Error('Network response was not ok')
 
@@ -351,6 +355,8 @@ export const handleStream = (
                 isFirstMessage = false
               } else if (bufferObj.event === 'agent_thought') {
                 onThought?.(bufferObj as ThoughtItem)
+              } else if (bufferObj.event === 'ui_part') {
+                onUIPart?.(bufferObj as UIPartStreamEvent)
               } else if (bufferObj.event === 'message_file') {
                 onFile?.(bufferObj as VisionFile)
               } else if (bufferObj.event === 'message_end') {
@@ -533,6 +539,7 @@ export const ssePost = async (
     onTTSEnd,
     onTextReplace,
     onAgentLog,
+    onUIPart,
     onError,
     getAbortController,
     onLoopStart,
@@ -660,6 +667,7 @@ export const ssePost = async (
         onDataSourceNodeError,
         onReasoning,
         onUnhandledEvent,
+        onUIPart,
       )
     })
     .catch((e) => {
@@ -702,6 +710,7 @@ export const sseGet = async (
     onTTSEnd,
     onTextReplace,
     onAgentLog,
+    onUIPart,
     onError,
     getAbortController,
     onLoopStart,
@@ -823,6 +832,7 @@ export const sseGet = async (
         onDataSourceNodeError,
         onReasoning,
         onUnhandledEvent,
+        onUIPart,
       )
     })
     .catch((e) => {

web/types/a2ui.ts

@@ -0,0 +1,156 @@
+export const A2UI_PROTOCOL_VERSION = 'v0.9.1' as const
+export const DIFY_A2UI_CATALOG_ID = 'https://dify.ai/a2ui/catalog/v1' as const
+
+export type JSONPrimitive = string | number | boolean | null
+export type JSONValue = JSONPrimitive | JSONValue[] | { [key: string]: JSONValue }
+
+export type A2UIBinding = {
+  path: string
+}
+
+export type A2UIDynamic<T extends JSONPrimitive> = T | A2UIBinding
+
+type A2UIComponentBase = {
+  id: string
+}
+
+export type A2UICard = A2UIComponentBase & {
+  component: 'Card'
+  children: string[]
+  title?: A2UIDynamic<string>
+}
+
+export type A2UIRow = A2UIComponentBase & {
+  component: 'Row'
+  children: string[]
+  gap?: 'small' | 'medium' | 'large'
+  align?: 'start' | 'center' | 'end'
+}
+
+export type A2UIColumn = A2UIComponentBase & {
+  component: 'Column'
+  children: string[]
+  gap?: 'small' | 'medium' | 'large'
+}
+
+export type A2UIText = A2UIComponentBase & {
+  component: 'Text'
+  text: A2UIDynamic<string>
+  variant?: 'body' | 'caption'
+}
+
+export type A2UIIcon = A2UIComponentBase & {
+  component: 'Icon'
+  name:
+    | 'clock'
+    | 'cloud'
+    | 'sun'
+    | 'rain'
+    | 'snow'
+    | 'wind'
+    | 'thermometer'
+    | 'calendar'
+    | 'location'
+}
+
+export type A2UIDivider = A2UIComponentBase & {
+  component: 'Divider'
+}
+
+export type A2UIBadge = A2UIComponentBase & {
+  component: 'Badge'
+  text: A2UIDynamic<string>
+  tone?: 'neutral' | 'info' | 'success' | 'warning' | 'critical'
+}
+
+export type A2UIMetric = A2UIComponentBase & {
+  component: 'Metric'
+  label: A2UIDynamic<string>
+  value: A2UIDynamic<JSONPrimitive>
+  unit?: A2UIDynamic<string>
+}
+
+export type A2UIDateTime = A2UIComponentBase & {
+  component: 'DateTime'
+  value: A2UIDynamic<string>
+  format?: 'date' | 'time' | 'datetime'
+}
+
+export type A2UIProgress = A2UIComponentBase & {
+  component: 'Progress'
+  value: A2UIDynamic<number>
+  max?: A2UIDynamic<number>
+  label: A2UIDynamic<string>
+}
+
+export type A2UIKeyValue = A2UIComponentBase & {
+  component: 'KeyValue'
+  label: A2UIDynamic<string>
+  value: A2UIDynamic<JSONPrimitive>
+}
+
+export type A2UIComponent =
+  | A2UICard
+  | A2UIRow
+  | A2UIColumn
+  | A2UIText
+  | A2UIIcon
+  | A2UIDivider
+  | A2UIBadge
+  | A2UIMetric
+  | A2UIDateTime
+  | A2UIProgress
+  | A2UIKeyValue
+
+type A2UIEnvelopeBase = {
+  version: typeof A2UI_PROTOCOL_VERSION
+}
+
+export type A2UICreateSurfaceEnvelope = A2UIEnvelopeBase & {
+  createSurface: {
+    surfaceId: string
+    catalogId: typeof DIFY_A2UI_CATALOG_ID
+  }
+}
+
+export type A2UIUpdateComponentsEnvelope = A2UIEnvelopeBase & {
+  updateComponents: {
+    surfaceId: string
+    components: A2UIComponent[]
+  }
+}
+
+export type A2UIUpdateDataModelEnvelope = A2UIEnvelopeBase & {
+  updateDataModel: {
+    surfaceId: string
+    path?: string
+    value: JSONValue
+  }
+}
+
+export type A2UIDeleteSurfaceEnvelope = A2UIEnvelopeBase & {
+  deleteSurface: {
+    surfaceId: string
+  }
+}
+
+export type A2UIEnvelope =
+  | A2UICreateSurfaceEnvelope
+  | A2UIUpdateComponentsEnvelope
+  | A2UIUpdateDataModelEnvelope
+  | A2UIDeleteSurfaceEnvelope
+
+export type UIPart = {
+  part_id: string
+  sequence: number
+  protocol: 'a2ui'
+  protocol_version: typeof A2UI_PROTOCOL_VERSION
+  messages: A2UIEnvelope[]
+  fallback?: string | null
+}
+
+export type UIPartStreamEvent = {
+  event: 'ui_part'
+  id: string
+  part: UIPart
+}