feat(inner_api): add endpoint to invalidate plugin model providers cache

Enterprise installs plugins into a workspace by talking to the plugin
daemon directly, so Dify never learns that a tenant's cached model
provider list is stale and serves it for the full 24h TTL.

Expose the existing PluginService.invalidate_plugin_model_providers_cache
over the inner API so enterprise can drop the cache for the tenants it
just changed. Batching is left to the caller.

api/configs/middleware/__init__.py

@@ -1,8 +1,17 @@
 import os
 from typing import Any, Literal, TypedDict, cast
 from urllib.parse import parse_qsl, quote_plus
+from urllib.parse import urlparse as _urlparse
 
-from pydantic import Field, NonNegativeFloat, NonNegativeInt, PositiveFloat, PositiveInt, computed_field
+from pydantic import (
+    Field,
+    NonNegativeFloat,
+    NonNegativeInt,
+    PositiveFloat,
+    PositiveInt,
+    computed_field,
+    model_validator,
+)
 from pydantic_settings import BaseSettings
 
 from .cache.redis_config import RedisConfig
@@ -408,4 +417,22 @@ class MiddlewareConfig(
     DatasetQueueMonitorConfig,
     MatrixoneConfig,
 ):
-    pass
+    @model_validator(mode="after")
+    def _validate_redis_urls_db_for_azure(self):
+        """Azure Managed Redis only supports db 0; reject non-zero db in Redis URLs."""
+        if not self.REDIS_USE_AZURE_MANAGED_IDENTITY:
+            return self
+
+        for url, name in (
+            (self.CELERY_BROKER_URL, "CELERY_BROKER_URL"),
+            (self.PUBSUB_REDIS_URL, "PUBSUB_REDIS_URL"),
+        ):
+            if not url:
+                continue
+            db: str = _urlparse(url).path.lstrip("/") or "0"
+            if db != "0":
+                raise ValueError(
+                    f"Azure Managed Redis only supports db 0, but {name} uses db {db}. "
+                    "Please set the db index to 0 in your URL."
+                )
+        return self

api/configs/middleware/cache/redis_config.py

@@ -1,4 +1,4 @@
-from pydantic import Field, NonNegativeInt, PositiveFloat, PositiveInt, field_validator
+from pydantic import Field, NonNegativeInt, PositiveFloat, PositiveInt, field_validator, model_validator
 from pydantic_settings import BaseSettings
 
 
@@ -42,6 +42,13 @@ class RedisConfig(BaseSettings):
         default=False,
     )
 
+    REDIS_USE_AZURE_MANAGED_IDENTITY: bool = Field(
+        description="Use Azure Managed Identity (Entra ID) for Redis authentication."
+        " When enabled, username/password are ignored and a token is acquired via DefaultAzureCredential."
+        " Requires azure-identity and redis-entraid packages.",
+        default=False,
+    )
+
     REDIS_SSL_CERT_REQS: str = Field(
         description="SSL certificate requirements (CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED)",
         default="CERT_NONE",
@@ -158,12 +165,33 @@ class RedisConfig(BaseSettings):
     REDIS_KEEPALIVE_INTERVAL: PositiveInt = Field(default=10, description="redis keepalive interval")
     REDIS_KEEPALIVE_COUNT: PositiveInt = Field(default=10, description="redis keepalive count")
 
-    @field_validator("REDIS_MAX_CONNECTIONS", mode="before")
+    @field_validator(
+        "REDIS_SSL_CA_CERTS",
+        "REDIS_SSL_CERTFILE",
+        "REDIS_SSL_KEYFILE",
+        "REDIS_MAX_CONNECTIONS",
+        mode="before",
+    )
     @classmethod
-    def _empty_string_to_none_for_max_conns(cls, v):
-        """Allow empty string in env/.env to mean 'unset' (None)."""
+    def _empty_string_to_none(cls, v):
+        """Allow empty string in env/.env to mean 'unset' (None).
+
+        Particularly important for SSL file paths: an empty string would cause
+        redis-py to call ``ssl.SSLContext.load_verify_locations(cafile="")``
+        which raises ``FileNotFoundError``.
+        """
         if v is None:
             return None
         if isinstance(v, str) and v.strip() == "":
             return None
         return v
+
+    @model_validator(mode="after")
+    def _validate_azure_managed_identity(self):
+        """Azure Managed Redis only supports db 0."""
+        if self.REDIS_USE_AZURE_MANAGED_IDENTITY and self.REDIS_DB != 0:
+            raise ValueError(
+                f"Azure Managed Redis only supports db 0, but REDIS_DB is set to {self.REDIS_DB}. "
+                "Please set REDIS_DB=0 when REDIS_USE_AZURE_MANAGED_IDENTITY is enabled."
+            )
+        return self

api/controllers/console/agent/roster.py

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

api/controllers/console/app/completion.py

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

api/controllers/console/auth/oauth_server.py

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

api/controllers/inner_api/__init__.py

@@ -23,6 +23,7 @@ from .knowledge import retrieval as _knowledge_retrieval
 from .plugin import agent_config as _agent_config
 from .plugin import agent_drive as _agent_drive
 from .plugin import plugin as _plugin
+from .workspace import plugin_model_providers as _plugin_model_providers
 from .workspace import workspace as _workspace
 
 api.add_namespace(inner_api_ns)
@@ -35,6 +36,7 @@ __all__ = [
     "_knowledge_retrieval",
     "_mail",
     "_plugin",
+    "_plugin_model_providers",
     "_runtime_credentials",
     "_workspace",
     "api",

api/controllers/inner_api/workspace/plugin_model_providers.py

@@ -0,0 +1,39 @@
+from flask_restx import Resource
+from pydantic import BaseModel, ConfigDict, Field
+
+from controllers.common.schema import register_schema_model
+from controllers.console.wraps import setup_required
+from controllers.inner_api import inner_api_ns
+from controllers.inner_api.wraps import enterprise_inner_api_only
+from core.plugin.plugin_service import PluginService
+
+
+class InvalidatePluginModelProvidersCachePayload(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    tenant_ids: list[str] = Field(default_factory=list, description="Workspace ids whose cache should be invalidated")
+
+
+register_schema_model(inner_api_ns, InvalidatePluginModelProvidersCachePayload)
+
+
+@inner_api_ns.route("/enterprise/workspace/plugin-model-providers/invalidate")
+class EnterprisePluginModelProvidersCacheInvalidate(Resource):
+    @setup_required
+    @enterprise_inner_api_only
+    @inner_api_ns.doc(
+        "enterprise_invalidate_plugin_model_providers_cache",
+        responses={
+            200: "Cache invalidated",
+            400: "Invalid request",
+            401: "Unauthorized - invalid API key",
+        },
+    )
+    @inner_api_ns.expect(inner_api_ns.models[InvalidatePluginModelProvidersCachePayload.__name__])
+    def post(self):
+        args = InvalidatePluginModelProvidersCachePayload.model_validate(inner_api_ns.payload or {})
+
+        for tenant_id in args.tenant_ids:
+            PluginService.invalidate_plugin_model_providers_cache(tenant_id)
+
+        return {"result": "success"}, 200

api/controllers/web/wraps.py

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

api/core/agent/base_agent_runner.py

@@ -10,13 +10,12 @@ 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, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager
 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
@@ -24,12 +23,6 @@ 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
@@ -359,34 +352,6 @@ 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,7 +235,6 @@ 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,
@@ -307,11 +306,9 @@ 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
@@ -334,7 +331,7 @@ class CotAgentRunner(BaseAgentRunner, ABC):
                 pass
 
         # invoke tool
-        invoke_result = ToolEngine.agent_invoke(
+        tool_invoke_response, message_files, tool_invoke_meta = ToolEngine.agent_invoke(
             session=session,
             tool=tool_instance,
             tool_parameters=tool_call_args,
@@ -345,13 +342,6 @@ 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
-                    invoke_result = ToolEngine.agent_invoke(
+                    tool_invoke_response, message_files, tool_invoke_meta = ToolEngine.agent_invoke(
                         session=session,
                         tool=tool_instance,
                         tool_parameters=tool_call_args,
@@ -265,13 +265,6 @@ 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,16 +52,8 @@ 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
@@ -443,29 +435,9 @@ class _AgentProcessRecorder:
         content = part.get("content")
         if content is None:
             content = part
-        thought_id = self._record_tool_observation(
-            tool_call_id=tool_call_id,
-            tool_name=tool_name,
-            observation=content,
-        )
+        self._record_tool_observation(tool_call_id=tool_call_id, tool_name=tool_name, observation=content)
 
-        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:
+    def _record_tool_observation(self, *, tool_call_id: str | None, tool_name: str | None, observation: Any) -> None:
         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:
@@ -473,34 +445,6 @@ 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,7 +7,6 @@ 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
@@ -44,7 +43,6 @@ 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"
@@ -459,13 +457,6 @@ 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,11 +2,10 @@ from collections.abc import Mapping, Sequence
 from enum import StrEnum
 from typing import Any, Literal
 
-from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
 
 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
@@ -31,14 +30,6 @@ 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):
@@ -82,7 +73,6 @@ 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"
@@ -202,14 +192,6 @@ 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,7 +26,6 @@ from core.app.entities.queue_entities import (
     QueuePingEvent,
     QueueRetrieverResourcesEvent,
     QueueStopEvent,
-    QueueUIPartEvent,
 )
 from core.app.entities.task_entities import (
     AgentMessageStreamResponse,
@@ -42,7 +41,6 @@ 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
@@ -54,7 +52,6 @@ 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
@@ -314,10 +311,6 @@ 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:
@@ -533,27 +526,6 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat
             task_id=self._application_generate_entity.task_id, id=message_id, answer=answer
         )
 
-    def _ui_part_to_stream_response(self, event: QueueUIPartEvent) -> UIPartStreamResponse | None:
-        """Upsert a UI part revision into persistent metadata and stream it."""
-        try:
-            updated_parts = upsert_ui_part(self._task_state.metadata.ui_parts, event.part)
-        except ValueError:
-            logger.warning(
-                "Ignored UI part that exceeds assistant message limits",
-                extra={"message_id": self._message_id, "part_id": event.part.part_id},
-                exc_info=True,
-            )
-            return None
-        if updated_parts is None:
-            return None
-        self._task_state.metadata.ui_parts = updated_parts
-
-        return UIPartStreamResponse(
-            task_id=self._application_generate_entity.task_id,
-            id=self._message_id,
-            part=event.part,
-        )
-
     def _agent_thought_to_stream_response(self, event: QueueAgentThoughtEvent) -> AgentThoughtStreamResponse | None:
         """
         Agent thought to stream response.

api/core/indexing_runner.py

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

api/core/logging/context.py

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

api/core/logging/filters.py

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

api/core/plugin/impl/base.py

@@ -1,7 +1,7 @@
 import inspect
 import json
 import logging
-from collections.abc import Callable, Generator, Mapping
+from collections.abc import Callable, Generator
 from typing import Any, cast
 from urllib.parse import unquote
 
@@ -23,7 +23,6 @@ from core.plugin.impl.exc import (
     PluginLLMPollingUnsupportedError,
     PluginNotFoundError,
     PluginPermissionDeniedError,
-    PluginRuntimeError,
     PluginUniqueIdentifierError,
 )
 from core.trigger.errors import (
@@ -376,18 +375,6 @@ class BasePluginClient:
                     # type `PluginLLMPollingUnsupportedError`.
                     case PluginLLMPollingUnsupportedError.__name__:
                         raise PluginLLMPollingUnsupportedError(description=error_object.get("message"))
-                    case PluginRuntimeError.__name__:
-                        args = error_object.get("args")
-                        lambda_request_id = args.get("request_id") if isinstance(args, Mapping) else None
-                        if not isinstance(lambda_request_id, str):
-                            lambda_request_id = None
-                        runtime_message = error_object.get("message")
-                        if not isinstance(runtime_message, str):
-                            runtime_message = "Plugin runtime request failed"
-                        raise PluginRuntimeError(
-                            description=runtime_message,
-                            lambda_request_id=lambda_request_id,
-                        )
                     case _:
                         raise PluginInvokeError(description=message)
             case PluginDaemonInternalServerError.__name__:

api/core/plugin/impl/exc.py

@@ -49,18 +49,6 @@ class PluginDaemonBadRequestError(PluginDaemonClientSideError):
     description: str = "Bad Request"
 
 
-class PluginRuntimeError(PluginDaemonInternalError):
-    """A plugin runtime failed before it could return a valid plugin response."""
-
-    lambda_request_id: str | None
-
-    def __init__(self, description: str, lambda_request_id: str | None = None) -> None:
-        self.lambda_request_id = lambda_request_id
-        if lambda_request_id:
-            description = description.replace(f"RequestId: {lambda_request_id} Error: ", "", 1)
-        super().__init__(description)
-
-
 class PluginInvokeError(PluginDaemonClientSideError, ValueError):
     description: str = "Invoke Error"
 

api/core/rag/docstore/dataset_docstore.py

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

api/core/rag/embedding/token_counter.py

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

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

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

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

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

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

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

api/core/tools/__base/tool.py

@@ -17,7 +17,6 @@ from core.tools.entities.tool_entities import (
     ToolParameter,
     ToolProviderType,
 )
-from core.tools.entities.ui_entities import ToolUIMessage
 
 
 class Tool(ABC):
@@ -296,13 +295,6 @@ 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,9 +7,6 @@ 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):
@@ -23,60 +20,19 @@ class CurrentTimeTool(BuiltinTool):
         app_id: str | None = None,
         message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
-        """Return the formatted time for the model and a standard time card for chat clients.
-
-        Invalid timezone values retain the existing text-only error response.
         """
-        timezone_name = tool_parameters.get("timezone", "UTC")
+        invoke tools
+        """
+        # get timezone
+        tz = 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)
+        if tz == "UTC":
+            yield self.create_text_message(f"{datetime.now(UTC).strftime(fm)}")
+            return
 
-        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",
-                                },
-                            ],
-                        },
-                    },
-                ],
-            }
-        )
+        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)}")

api/core/tools/entities/tool_entities.py

@@ -32,7 +32,6 @@ 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):
@@ -241,15 +240,13 @@ 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: (
-        ToolUIMessage
-        | JsonMessage
+        JsonMessage
         | TextMessage
         | BlobChunkMessage
         | BlobMessage
@@ -278,8 +275,6 @@ 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

@@ -1,675 +0,0 @@
-"""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,7 +3,6 @@ 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
@@ -22,12 +21,6 @@ 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,
@@ -47,16 +40,6 @@ 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.
@@ -76,13 +59,9 @@ class ToolEngine:
         conversation_id: str | None = None,
         app_id: str | None = None,
         message_id: str | None = None,
-    ) -> ToolAgentInvokeResult:
+    ) -> tuple[str, list[str], ToolInvokeMeta]:
         """
-        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.
+        Agent invokes the tool with the given arguments.
         """
         # check if arguments is a string
         if isinstance(tool_parameters, str):
@@ -124,7 +103,7 @@ class ToolEngine:
                 conversation_id=message.conversation_id,
             )
 
-            message_list, ui_messages = ToolEngine.collect_agent_messages(messages)
+            message_list = list(messages)
 
             # extract binary data from tool invoke message
             binary_files = ToolEngine._extract_tool_response_binary_and_text(message_list)
@@ -147,12 +126,7 @@ class ToolEngine:
             )
 
             # transform tool invoke message to get LLM friendly message
-            return ToolAgentInvokeResult(
-                observation=plain_text,
-                message_files=message_files,
-                ui_messages=ui_messages,
-                meta=meta,
-            )
+            return plain_text, message_files, meta
         except ToolProviderCredentialValidationError as e:
             logger.error(e, exc_info=True)
             error_response = "Please check your tool provider credentials"
@@ -174,23 +148,13 @@ class ToolEngine:
             error_response = f"tool invoke error: {meta.error}"
             agent_tool_callback.on_tool_error(e)
             logger.error(e, exc_info=True)
-            return ToolAgentInvokeResult(
-                observation=error_response,
-                message_files=[],
-                ui_messages=[],
-                meta=meta,
-            )
+            return error_response, [], 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 ToolAgentInvokeResult(
-            observation=error_response,
-            message_files=[],
-            ui_messages=[],
-            meta=ToolInvokeMeta.error_instance(error_response),
-        )
+        return error_response, [], ToolInvokeMeta.error_instance(error_response)
 
     @staticmethod
     def generic_invoke(
@@ -233,7 +197,7 @@ class ToolEngine:
                 tool_outputs=response,
             )
 
-            return ToolEngine.normalize_ui_messages(response)
+            return response
         except Exception as e:
             workflow_tool_callback.on_tool_error(e)
             raise e
@@ -302,10 +266,7 @@ class ToolEngine:
                         ensure_ascii=False,
                     )
                 )
-            elif response.type in {
-                ToolInvokeMessage.MessageType.VARIABLE,
-                ToolInvokeMessage.MessageType.UI,
-            }:
+            elif response.type == ToolInvokeMessage.MessageType.VARIABLE:
                 continue
             else:
                 parts.append(str(response.message))
@@ -317,124 +278,6 @@ 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,7 +90,6 @@ 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
 
@@ -667,11 +666,6 @@ 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
@@ -692,7 +686,6 @@ class DifyToolNodeRuntime(ToolNodeRuntimeProtocol):
         | CoreToolInvokeMessage.FileMessage
         | CoreToolInvokeMessage.VariableMessage
         | CoreToolInvokeMessage.RetrieverResourceMessage
-        | ToolUIMessage
         | None,
     ) -> (
         ToolRuntimeMessage.TextMessage

api/extensions/azure.py

@@ -0,0 +1,58 @@
+"""Azure-specific helpers for Redis authentication via Entra ID (Managed Identity)."""
+
+from typing import Union, override
+
+from redis import CredentialProvider
+
+AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
+
+
+class AzureEntraIdCredentialProvider(CredentialProvider):
+    """Redis credential provider for Azure Entra ID (Managed Identity) authentication.
+
+    Wraps ``redis-entraid``'s provider so that it can be instantiated with no
+    arguments — required by kombu's URL-based ``credential_provider`` resolution.
+    """
+
+    _inner: CredentialProvider
+
+    def __init__(self) -> None:
+        from redis_entraid.cred_provider import create_from_default_azure_credential
+
+        self._inner = create_from_default_azure_credential(
+            scopes=(AZURE_REDIS_SCOPE,),
+        )
+
+    @override
+    def get_credentials(self) -> Union[tuple[str], tuple[str, str]]:
+        return self._inner.get_credentials()
+
+
+def get_azure_credential_provider() -> CredentialProvider:
+    """Create a redis-py credential provider for Azure Entra ID authentication."""
+    from redis_entraid.cred_provider import create_from_default_azure_credential
+
+    return create_from_default_azure_credential(
+        scopes=(AZURE_REDIS_SCOPE,),
+    )
+
+
+def apply_azure_redis_auth(params: dict) -> None:
+    """Apply Azure Entra ID authentication to a Redis connection params dict.
+
+    Removes static username/password and injects a credential_provider instead.
+    """
+    params.pop("username", None)
+    params.pop("password", None)
+    params["credential_provider"] = get_azure_credential_provider()
+
+
+def apply_azure_celery_broker_auth(celery_app, broker_url: str) -> None:
+    """Configure Celery broker to authenticate via Azure Entra ID credential provider."""
+    cred_param = "credential_provider=extensions.azure.AzureEntraIdCredentialProvider"
+    sep = "&" if "?" in broker_url else "?"
+    broker_url_with_cred = f"{broker_url}{sep}{cred_param}"
+    celery_app.conf.update(
+        broker_read_url=broker_url_with_cred,
+        broker_write_url=broker_url_with_cred,
+    )

api/extensions/ext_celery.py

@@ -1,3 +1,4 @@
+import logging
 import ssl
 from datetime import timedelta
 from typing import Any
@@ -9,9 +10,12 @@ from typing_extensions import TypedDict
 
 from configs import dify_config
 from dify_app import DifyApp
+from extensions.azure import AzureEntraIdCredentialProvider, apply_azure_celery_broker_auth
 from extensions.redis_names import normalize_redis_key_prefix
 from extensions.workflow_warm_shutdown import setup_workflow_warm_shutdown_handler
 
+logger = logging.getLogger(__name__)
+
 
 class _CelerySentinelKwargsDict(TypedDict):
     socket_timeout: float | None
@@ -108,10 +112,12 @@ def init_app(app: DifyApp) -> Celery:
 
     broker_transport_options = get_celery_broker_transport_options()
 
+    broker_url = dify_config.CELERY_BROKER_URL
+
     celery_app = Celery(
         app.name,
         task_cls=FlaskTask,
-        broker=dify_config.CELERY_BROKER_URL,
+        broker=broker_url,
         backend=dify_config.CELERY_BACKEND,
     )
 
@@ -128,9 +134,12 @@ def init_app(app: DifyApp) -> Celery:
     )
 
     if dify_config.CELERY_BACKEND == "redis":
-        celery_app.conf.update(
-            result_backend_transport_options=broker_transport_options,
-        )
+        redis_backend_conf: dict[str, Any] = {
+            "result_backend_transport_options": broker_transport_options,
+        }
+        if dify_config.REDIS_USE_AZURE_MANAGED_IDENTITY:
+            redis_backend_conf["redis_backend_credential_provider"] = AzureEntraIdCredentialProvider()
+        celery_app.conf.update(**redis_backend_conf)
 
     # Apply SSL configuration if enabled
     ssl_options = get_celery_ssl_options()
@@ -141,6 +150,10 @@ def init_app(app: DifyApp) -> Celery:
             redis_backend_use_ssl=ssl_options if dify_config.CELERY_BACKEND == "redis" else None,
         )
 
+    if dify_config.REDIS_USE_AZURE_MANAGED_IDENTITY and broker_url:
+        apply_azure_celery_broker_auth(celery_app, broker_url)
+        logger.info("Celery broker: using Azure Managed Identity (Entra ID) authentication")
+
     if dify_config.LOG_FILE:
         celery_app.conf.update(
             worker_logfile=dify_config.LOG_FILE,

api/extensions/ext_login.py

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

api/extensions/ext_redis.py

@@ -21,6 +21,7 @@ from typing_extensions import TypedDict
 
 from configs import dify_config
 from dify_app import DifyApp
+from extensions.azure import apply_azure_redis_auth, get_azure_credential_provider
 from extensions.redis_names import (
     normalize_redis_key_prefix,
     serialize_redis_name,
@@ -438,6 +439,10 @@ def _create_standalone_client(redis_params: RedisBaseParamsDict) -> Union[redis.
         "connection_class": connection_class,
     }
 
+    if dify_config.REDIS_USE_AZURE_MANAGED_IDENTITY:
+        apply_azure_redis_auth(params)
+        logger.info("Redis: using Azure Managed Identity (Entra ID) authentication")
+
     if dify_config.REDIS_MAX_CONNECTIONS:
         params["max_connections"] = dify_config.REDIS_MAX_CONNECTIONS
 
@@ -457,12 +462,20 @@ def _create_pubsub_client(pubsub_url: str, use_clusters: bool) -> redis.Redis |
         kwargs: dict[str, Any] = {**health_params}
         if max_conns:
             kwargs["max_connections"] = max_conns
+        if dify_config.REDIS_USE_AZURE_MANAGED_IDENTITY:
+            kwargs["credential_provider"] = get_azure_credential_provider()
+            kwargs["ssl_cert_reqs"] = ssl.CERT_NONE
+            logger.info("PubSub Redis (cluster): using Azure Managed Identity (Entra ID) authentication")
         return RedisCluster.from_url(pubsub_url, **kwargs)
 
     standalone_health_params: dict[str, Any] = dict(_get_connection_health_params())
     kwargs = {**standalone_health_params}
     if max_conns:
         kwargs["max_connections"] = max_conns
+    if dify_config.REDIS_USE_AZURE_MANAGED_IDENTITY:
+        kwargs["credential_provider"] = get_azure_credential_provider()
+        kwargs["ssl_cert_reqs"] = ssl.CERT_NONE
+        logger.info("PubSub Redis: using Azure Managed Identity (Entra ID) authentication")
     return redis.Redis.from_url(pubsub_url, **kwargs)
 
 
@@ -486,7 +499,7 @@ def init_app(app: DifyApp):
 
     global _pubsub_redis_client
     _pubsub_redis_client = client
-    if dify_config.normalized_pubsub_redis_url:
+    if dify_config.PUBSUB_REDIS_URL:
         _pubsub_redis_client = _create_pubsub_client(
             dify_config.normalized_pubsub_redis_url, dify_config.PUBSUB_REDIS_USE_CLUSTERS
         )

api/extensions/otel/runtime.py

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

api/fields/message_fields.py

@@ -7,7 +7,6 @@ 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
@@ -65,7 +64,6 @@ 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:
@@ -81,12 +79,6 @@ 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/libs/external_api.py

@@ -9,8 +9,6 @@ from werkzeug.http import HTTP_STATUS_CODES
 
 from configs import dify_config
 from core.errors.error import AppInvokeQuotaExceededError
-from core.plugin.impl.exc import PluginRuntimeError
-from extensions.ext_logging import get_request_id
 from libs.flask_restx_compat import install_swagger_compatibility
 from libs.token import build_force_logout_cookie_headers
 
@@ -102,20 +100,6 @@ def register_external_error_handlers(api: Api, body_formatter: ErrorBodyFormatte
         data = {"code": "too_many_requests", "message": str(e), "status": status_code}
         return _finalize(e, data, status_code), status_code
 
-    def handle_plugin_runtime_error(e: PluginRuntimeError):
-        got_request_exception.send(current_app, exception=e)
-        status_code = 502
-        details = {"request_id": get_request_id()}
-        if e.lambda_request_id:
-            details["lambda_request_id"] = e.lambda_request_id
-        data = {
-            "code": "plugin_runtime_error",
-            "message": e.description,
-            "details": details,
-            "status": status_code,
-        }
-        return _finalize(e, data, status_code), status_code
-
     def handle_general_exception(e: Exception):
         got_request_exception.send(current_app, exception=e)
 
@@ -137,7 +121,6 @@ def register_external_error_handlers(api: Api, body_formatter: ErrorBodyFormatte
     api.errorhandler(HTTPException)(handle_http_exception)
     api.errorhandler(ValueError)(handle_value_error)
     api.errorhandler(AppInvokeQuotaExceededError)(handle_quota_exceeded)
-    api.errorhandler(PluginRuntimeError)(handle_plugin_runtime_error)
     api.errorhandler(Exception)(handle_general_exception)
 
 

api/openapi/markdown/console-openapi.md

@@ -19559,7 +19559,6 @@ Coarse node-level status used by Inspector to pick a banner.
 | ---- | ---- | ----------- | -------- |
 | avatar | string |  | No |
 | email | string |  | Yes |
-| id | string |  | Yes |
 | interface_language | string |  | Yes |
 | name | string |  | Yes |
 | timezone | string |  | Yes |

api/pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "psycopg2-binary>=2.9.12,<3.0.0",
     "python-socketio>=5.13.0,<6.0.0",
     "redis[hiredis]>=7.4.0,<8.0.0",
+    "redis-entraid>=1.2.0,<2.0.0",
     "sendgrid>=6.12.5,<7.0.0",
     "sseclient-py>=1.8.0,<2.0.0",
     # Stable: production-proven, cap below the next major

api/services/agent/observability_service.py

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

api/services/agent/roster_service.py

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

api/services/agent_tool_inner_service.py

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

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

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

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

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

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

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

api/tests/unit_tests/controllers/inner_api/workspace/test_plugin_model_providers.py

@@ -0,0 +1,64 @@
+import inspect
+from unittest.mock import call, patch
+
+import pytest
+from flask import Flask
+from pydantic import ValidationError
+
+from controllers.inner_api.workspace.plugin_model_providers import (
+    EnterprisePluginModelProvidersCacheInvalidate,
+    InvalidatePluginModelProvidersCachePayload,
+)
+
+
+class TestInvalidatePluginModelProvidersCachePayload:
+    def test_valid_payload(self):
+        payload = InvalidatePluginModelProvidersCachePayload.model_validate(
+            {"tenant_ids": ["tenant-alpha", "tenant-beta"]}
+        )
+        assert payload.tenant_ids == ["tenant-alpha", "tenant-beta"]
+
+    def test_missing_tenant_ids_defaults_to_empty(self):
+        payload = InvalidatePluginModelProvidersCachePayload.model_validate({})
+        assert payload.tenant_ids == []
+
+    def test_unknown_field_rejected(self):
+        with pytest.raises(ValidationError):
+            InvalidatePluginModelProvidersCachePayload.model_validate({"tenant_ids": ["tenant-alpha"], "generation": 7})
+
+
+class TestEnterprisePluginModelProvidersCacheInvalidate:
+    @pytest.fixture
+    def api_instance(self):
+        return EnterprisePluginModelProvidersCacheInvalidate()
+
+    def _post(self, api_instance, app: Flask, payload):
+        unwrapped_post = inspect.unwrap(api_instance.post)
+        with app.test_request_context():
+            with patch("controllers.inner_api.workspace.plugin_model_providers.inner_api_ns") as mock_ns:
+                mock_ns.payload = payload
+                return unwrapped_post(api_instance)
+
+    @patch("controllers.inner_api.workspace.plugin_model_providers.PluginService")
+    def test_post_invalidates_once_per_tenant(self, mock_plugin_service, api_instance, app: Flask):
+        result = self._post(api_instance, app, {"tenant_ids": ["tenant-alpha", "tenant-beta"]})
+
+        assert result == ({"result": "success"}, 200)
+        assert mock_plugin_service.invalidate_plugin_model_providers_cache.call_args_list == [
+            call("tenant-alpha"),
+            call("tenant-beta"),
+        ]
+
+    @patch("controllers.inner_api.workspace.plugin_model_providers.PluginService")
+    def test_post_with_empty_list_is_a_no_op(self, mock_plugin_service, api_instance, app: Flask):
+        result = self._post(api_instance, app, {"tenant_ids": []})
+
+        assert result == ({"result": "success"}, 200)
+        mock_plugin_service.invalidate_plugin_model_providers_cache.assert_not_called()
+
+    @patch("controllers.inner_api.workspace.plugin_model_providers.PluginService")
+    def test_post_with_missing_payload_is_a_no_op(self, mock_plugin_service, api_instance, app: Flask):
+        result = self._post(api_instance, app, None)
+
+        assert result == ({"result": "success"}, 200)
+        mock_plugin_service.invalidate_plugin_model_providers_cache.assert_not_called()

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

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

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

@@ -8,8 +8,6 @@ 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
@@ -39,45 +37,6 @@ 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,7 +10,6 @@ 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
 
 
@@ -29,19 +28,6 @@ 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
@@ -195,7 +181,7 @@ class TestHandleInvokeAction:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=_invoke_result(observation="result"),
+            return_value=("result", [], MagicMock(to_dict=lambda: {})),
         )
 
         response, meta = runner._handle_invoke_action(runner.session, action, tool_instances, [])
@@ -240,7 +226,7 @@ class TestRun:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=_invoke_result(),
+            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
         )
 
         runner.agent_callback = None
@@ -262,7 +248,7 @@ class TestRun:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=_invoke_result(),
+            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
         )
 
         runner.agent_callback = None
@@ -337,7 +323,7 @@ class TestRun:
         runner.model_instance.invoke_llm = MagicMock(return_value=[])
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=_invoke_result(),
+            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
         )
 
         fake_prompt_tool = MagicMock()
@@ -405,7 +391,7 @@ class TestRun:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=_invoke_result(),
+            return_value=("ok", [], MagicMock(to_dict=lambda: {})),
         )
 
         runner.app_config.agent.max_iteration = 5
@@ -466,10 +452,7 @@ class TestHandleInvokeActionExtended:
 
         mocker.patch(
             "core.agent.cot_agent_runner.ToolEngine.agent_invoke",
-            return_value=_invoke_result(
-                message_files=["file1"],
-                meta=MagicMock(to_dict=lambda: {"k": "v"}),
-            ),
+            return_value=("ok", ["file1"], MagicMock(to_dict=lambda: {"k": "v"})),
         )
 
         message_file_ids = []

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

@@ -12,7 +12,6 @@ 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,
@@ -440,12 +439,7 @@ class TestRunMethod:
         tool_invoke_meta.to_dict.return_value = {"ok": True}
         mocker.patch(
             "core.agent.fc_agent_runner.ToolEngine.agent_invoke",
-            return_value=ToolAgentInvokeResult(
-                observation="ok",
-                message_files=["file1"],
-                ui_messages=[],
-                meta=tool_invoke_meta,
-            ),
+            return_value=("ok", ["file1"], 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,9 +55,7 @@ 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
@@ -1007,150 +1005,6 @@ 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

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

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

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

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

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

api/tests/unit_tests/core/plugin/impl/test_base_client_impl.py

@@ -7,7 +7,7 @@ from pytest_mock import MockerFixture
 from core.plugin.endpoint.exc import EndpointSetupFailedError
 from core.plugin.entities.plugin_daemon import PluginDaemonInnerError
 from core.plugin.impl.base import PLUGIN_DAEMON_MAX_PATH_LENGTH, BasePluginClient
-from core.plugin.impl.exc import PluginLLMPollingUnsupportedError, PluginRuntimeError
+from core.plugin.impl.exc import PluginLLMPollingUnsupportedError
 from core.trigger.errors import (
     EventIgnoreError,
     TriggerInvokeError,
@@ -175,25 +175,3 @@ class TestBasePluginClientImpl:
 
         with pytest.raises(PluginLLMPollingUnsupportedError):
             client._handle_plugin_daemon_error("PluginInvokeError", message)
-
-    def test_handle_plugin_daemon_error_maps_runtime_error_to_typed_exception(self):
-        client = BasePluginClient()
-        lambda_request_id = "45664803-3d3c-4d4f-93fe-e3b19e43092b"
-        message = json.dumps(
-            {
-                "error_type": PluginRuntimeError.__name__,
-                "message": (
-                    "Plugin runtime request failed: Runtime.ExitError: "
-                    f"RequestId: {lambda_request_id} Error: Runtime exited with error: exit status 1"
-                ),
-                "args": {"request_id": lambda_request_id, "status_code": 200},
-            }
-        )
-
-        with pytest.raises(PluginRuntimeError) as exc_info:
-            client._handle_plugin_daemon_error("PluginInvokeError", message)
-
-        assert exc_info.value.description == (
-            "Plugin runtime request failed: Runtime.ExitError: Runtime exited with error: exit status 1"
-        )
-        assert exc_info.value.lambda_request_id == lambda_request_id

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -18,7 +18,6 @@ from core.tools.entities.tool_entities import (
     ToolParameter,
     ToolProviderType,
 )
-from core.tools.entities.ui_entities import A2UI_CATALOG_ID
 
 
 class DummyCastType:
@@ -366,29 +365,6 @@ 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 UTC, date, datetime
+from datetime import date
 from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 from zoneinfo import ZoneInfo
 
 import pytest
@@ -28,7 +28,6 @@ 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
@@ -54,67 +53,10 @@ def _raise_runtime_error(*_args: object, **_kwargs: object) -> None:
 
 def test_current_time_tool():
     current_tool = _build_builtin_tool(CurrentTimeTool)
-    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",
-                        },
-                    ],
-                },
-            },
-        ],
-    }
+    utc_text = list(current_tool.invoke(session=MagicMock(), user_id="u", tool_parameters={"timezone": "UTC"}))[
+        0
+    ].message.text
+    assert utc_text
 
     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,7 +20,6 @@ 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,
@@ -211,7 +210,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 = ToolEngine.agent_invoke(
+                    result_text, message_files, result_meta = ToolEngine.agent_invoke(
                         session=MagicMock(),
                         tool=tool,
                         tool_parameters="hello",
@@ -222,158 +221,20 @@ def test_agent_invoke_success():
                         agent_tool_callback=callback,
                     )
 
-    assert result.observation == "ok"
-    assert result.message_files == []
-    assert result.ui_messages == []
-    assert result.meta.error is None
+    assert result_text == "ok"
+    assert message_files == []
+    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")):
-        result = ToolEngine.agent_invoke(
+        error_text, files, error_meta = ToolEngine.agent_invoke(
             session=MagicMock(),
             tool=tool,
             tool_parameters={"a": 1},
@@ -384,9 +245,9 @@ def test_agent_invoke_param_validation_error():
             agent_tool_callback=callback,
         )
 
-    assert "tool parameters validation error" in result.observation
-    assert result.message_files == []
-    assert result.meta.error
+    assert "tool parameters validation error" in error_text
+    assert files == []
+    assert error_meta.error
 
 
 def test_agent_invoke_engine_meta_error():
@@ -396,7 +257,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):
-        result = ToolEngine.agent_invoke(
+        error_text, files, error_meta = ToolEngine.agent_invoke(
             session=MagicMock(),
             tool=tool,
             tool_parameters={"a": 1},
@@ -407,9 +268,9 @@ def test_agent_invoke_engine_meta_error():
             agent_tool_callback=callback,
         )
 
-    assert "meta failure" in result.observation
-    assert result.message_files == []
-    assert result.meta.error == "meta failure"
+    assert "meta failure" in error_text
+    assert files == []
+    assert error_meta.error == "meta failure"
 
 
 def test_convert_tool_response_excludes_variable_messages():
@@ -440,7 +301,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")):
-        result = ToolEngine.agent_invoke(
+        error_text, files, _ = ToolEngine.agent_invoke(
             session=MagicMock(),
             tool=tool,
             tool_parameters={"a": 1},
@@ -451,5 +312,5 @@ def test_agent_invoke_tool_invoke_error():
             agent_tool_callback=callback,
         )
 
-    assert "tool invoke error" in result.observation
-    assert result.message_files == []
+    assert "tool invoke error" in error_text
+    assert files == []

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

@@ -1,500 +0,0 @@
-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,7 +13,6 @@ 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
@@ -328,29 +327,6 @@ def test_convert_message_payload_rejects_unknown_types(runtime: DifyToolNodeRunt
         runtime._convert_message_payload(object())
 
 
-def test_adapt_messages_ignores_chat_ui_messages(runtime: DifyToolNodeRuntime) -> None:
-    ui_message = ToolInvokeMessage(
-        type=ToolInvokeMessage.MessageType.UI,
-        message={
-            "messages": [
-                {
-                    "version": "v0.9.1",
-                    "createSurface": {"surfaceId": "clock", "catalogId": A2UI_CATALOG_ID},
-                },
-                {
-                    "version": "v0.9.1",
-                    "updateComponents": {
-                        "surfaceId": "clock",
-                        "components": [{"id": "root", "component": "Text", "text": "10:30"}],
-                    },
-                },
-            ]
-        },
-    )
-
-    assert list(runtime._adapt_messages(iter([ui_message]), provider_name="provider")) == []
-
-
 def test_resolve_provider_icons_prefers_builtin_tool_icons(runtime: DifyToolNodeRuntime) -> None:
     plugin = SimpleNamespace(
         plugin_id="langgenius/tools",

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

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

api/tests/unit_tests/extensions/test_azure.py

@@ -0,0 +1,186 @@
+"""Tests for Azure Managed Identity Redis helpers."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from configs import DifyConfig
+from configs.middleware.cache.redis_config import RedisConfig
+from extensions.azure import (
+    AzureEntraIdCredentialProvider,
+    apply_azure_celery_broker_auth,
+    apply_azure_redis_auth,
+    get_azure_credential_provider,
+)
+
+
+class TestAzureConfigValidation:
+    """Test pydantic config validation for Azure Managed Redis constraints."""
+
+    def test_redis_db_0_passes(self):
+        config = RedisConfig(REDIS_USE_AZURE_MANAGED_IDENTITY=True, REDIS_DB=0)
+        assert config.REDIS_DB == 0
+
+    def test_redis_db_nonzero_raises(self):
+        with pytest.raises(ValueError, match="only supports db 0"):
+            RedisConfig(REDIS_USE_AZURE_MANAGED_IDENTITY=True, REDIS_DB=1)
+
+    def test_redis_db_nonzero_without_azure_mi_passes(self):
+        config = RedisConfig(REDIS_USE_AZURE_MANAGED_IDENTITY=False, REDIS_DB=5)
+        assert config.REDIS_DB == 5
+
+    def test_redis_db_default_with_azure_mi_passes(self):
+        config = RedisConfig(REDIS_USE_AZURE_MANAGED_IDENTITY=True)
+        assert config.REDIS_DB == 0
+
+    def test_celery_broker_url_db_0_with_azure_mi_passes(self):
+        config = DifyConfig(
+            REDIS_USE_AZURE_MANAGED_IDENTITY=True,
+            REDIS_DB=0,
+            CELERY_BROKER_URL="rediss://:@host:10000/0",
+        )
+        assert config.CELERY_BROKER_URL == "rediss://:@host:10000/0"
+
+    def test_celery_broker_url_nonzero_db_with_azure_mi_raises(self):
+        with pytest.raises(ValueError, match="only supports db 0"):
+            DifyConfig(
+                REDIS_USE_AZURE_MANAGED_IDENTITY=True,
+                REDIS_DB=0,
+                CELERY_BROKER_URL="rediss://:@host:10000/1",
+            )
+
+    def test_celery_broker_url_nonzero_db_without_azure_mi_passes(self):
+        config = DifyConfig(
+            REDIS_USE_AZURE_MANAGED_IDENTITY=False,
+            CELERY_BROKER_URL="redis://localhost:6379/5",
+        )
+        assert config.CELERY_BROKER_URL == "redis://localhost:6379/5"
+
+    def test_pubsub_redis_url_db_0_with_azure_mi_passes(self):
+        config = DifyConfig(
+            REDIS_USE_AZURE_MANAGED_IDENTITY=True,
+            REDIS_DB=0,
+            PUBSUB_REDIS_URL="rediss://:@host:10000/0",
+        )
+        assert config.PUBSUB_REDIS_URL == "rediss://:@host:10000/0"
+
+    def test_pubsub_redis_url_nonzero_db_with_azure_mi_raises(self):
+        with pytest.raises(ValueError, match="only supports db 0"):
+            DifyConfig(
+                REDIS_USE_AZURE_MANAGED_IDENTITY=True,
+                REDIS_DB=0,
+                PUBSUB_REDIS_URL="rediss://:@host:10000/2",
+            )
+
+    def test_pubsub_redis_url_nonzero_db_without_azure_mi_passes(self):
+        config = DifyConfig(
+            REDIS_USE_AZURE_MANAGED_IDENTITY=False,
+            PUBSUB_REDIS_URL="redis://localhost:6379/3",
+        )
+        assert config.PUBSUB_REDIS_URL == "redis://localhost:6379/3"
+
+
+class TestApplyAzureCeleryBrokerAuth:
+    """Test apply_azure_celery_broker_auth Celery configuration."""
+
+    def test_sets_broker_read_and_write_urls(self):
+        mock_app = MagicMock()
+        apply_azure_celery_broker_auth(mock_app, "rediss://:@host:10000/0")
+
+        mock_app.conf.update.assert_called_once()
+        call_kwargs = mock_app.conf.update.call_args[1]
+
+        assert "broker_read_url" in call_kwargs
+        assert "broker_write_url" in call_kwargs
+
+    def test_appends_credential_provider_query_param(self):
+        mock_app = MagicMock()
+        apply_azure_celery_broker_auth(mock_app, "rediss://:@host:10000/0")
+
+        call_kwargs = mock_app.conf.update.call_args[1]
+        expected_param = "credential_provider=extensions.azure.AzureEntraIdCredentialProvider"
+        assert expected_param in call_kwargs["broker_read_url"]
+        assert expected_param in call_kwargs["broker_write_url"]
+
+    def test_uses_ampersand_when_url_already_has_query(self):
+        mock_app = MagicMock()
+        apply_azure_celery_broker_auth(mock_app, "rediss://:@host:10000/0?timeout=5")
+
+        call_kwargs = mock_app.conf.update.call_args[1]
+        assert "&credential_provider=" in call_kwargs["broker_read_url"]
+
+    def test_uses_question_mark_when_url_has_no_query(self):
+        mock_app = MagicMock()
+        apply_azure_celery_broker_auth(mock_app, "rediss://:@host:10000/0")
+
+        call_kwargs = mock_app.conf.update.call_args[1]
+        assert "?credential_provider=" in call_kwargs["broker_read_url"]
+
+
+class TestAzureEntraIdCredentialProvider:
+    """Test AzureEntraIdCredentialProvider wrapper."""
+
+    def test_get_credentials_delegates_to_inner(self):
+        provider = AzureEntraIdCredentialProvider.__new__(AzureEntraIdCredentialProvider)
+        mock_inner = MagicMock()
+        mock_inner.get_credentials.return_value = ("user-oid", "jwt-token")
+        provider._inner = mock_inner
+
+        result = provider.get_credentials()
+        assert result == ("user-oid", "jwt-token")
+
+
+class TestGetAzureCredentialProvider:
+    """Test get_azure_credential_provider factory."""
+
+    @patch("redis_entraid.cred_provider.create_from_default_azure_credential")
+    def test_calls_create_with_correct_scope(self, mock_create):
+        mock_create.return_value = MagicMock()
+        get_azure_credential_provider()
+
+        mock_create.assert_called_once_with(
+            scopes=("https://redis.azure.com/.default",),
+        )
+
+    @patch("redis_entraid.cred_provider.create_from_default_azure_credential")
+    def test_returns_provider_instance(self, mock_create):
+        sentinel = MagicMock()
+        mock_create.return_value = sentinel
+
+        result = get_azure_credential_provider()
+        assert result is sentinel
+
+
+class TestApplyAzureRedisAuth:
+    """Test apply_azure_redis_auth params mutation."""
+
+    @patch("extensions.azure.get_azure_credential_provider")
+    def test_removes_username_and_password(self, mock_get_provider):
+        mock_get_provider.return_value = MagicMock()
+        params: dict = {"host": "localhost", "username": "user", "password": "secret"}
+
+        apply_azure_redis_auth(params)
+
+        assert "username" not in params
+        assert "password" not in params
+
+    @patch("extensions.azure.get_azure_credential_provider")
+    def test_injects_credential_provider(self, mock_get_provider):
+        sentinel = MagicMock()
+        mock_get_provider.return_value = sentinel
+        params: dict = {"host": "localhost", "username": "u", "password": "p"}
+
+        apply_azure_redis_auth(params)
+
+        assert params["credential_provider"] is sentinel
+
+    @patch("extensions.azure.get_azure_credential_provider")
+    def test_handles_missing_username_password(self, mock_get_provider):
+        mock_get_provider.return_value = MagicMock()
+        params: dict = {"host": "localhost"}
+
+        apply_azure_redis_auth(params)
+
+        assert "username" not in params
+        assert "password" not in params
+        assert "credential_provider" in params

api/tests/unit_tests/extensions/test_ext_login.py

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

api/tests/unit_tests/fields/test_message_fields.py

@@ -1,6 +1,5 @@
 from decimal import Decimal
 
-from core.tools.entities.ui_entities import A2UI_CATALOG_ID, MessageUIPart, ToolUIMessage
 from fields.message_fields import ExploreMessageListItem, MessageListItem, WebMessageListItem
 
 
@@ -19,44 +18,6 @@ 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)
@@ -105,46 +66,3 @@ 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/libs/test_external_api.py

@@ -4,7 +4,6 @@ from werkzeug.exceptions import BadRequest, Unauthorized
 
 from constants import COOKIE_NAME_ACCESS_TOKEN, COOKIE_NAME_CSRF_TOKEN, COOKIE_NAME_REFRESH_TOKEN
 from core.errors.error import AppInvokeQuotaExceededError
-from core.plugin.impl.exc import PluginRuntimeError
 from libs.exception import BaseHTTPException
 from libs.external_api import ExternalApi
 from libs.rate_limit import _BearerRateLimited
@@ -40,14 +39,6 @@ def _create_api_app():
         def get(self):
             raise RuntimeError("oops")
 
-    @api.route("/plugin-runtime-error")
-    class PluginRuntime(Resource):
-        def get(self):
-            raise PluginRuntimeError(
-                "Plugin runtime request failed: Runtime.ExitError: Runtime exited with error: exit status 1",
-                lambda_request_id="lambda-request-id",
-            )
-
     # Note: We avoid altering default_mediatype to keep normal error paths
 
     # Special 400 message rewrite
@@ -116,24 +107,6 @@ def test_external_api_json_message_and_bad_request_rewrite():
     assert res.get_json()["message"] == "Invalid JSON payload received or JSON payload is empty."
 
 
-def test_external_api_plugin_runtime_error(mocker):
-    mocker.patch("libs.external_api.get_request_id", return_value="api-request-id")
-    app = _create_api_app()
-
-    res = app.test_client().get("/api/plugin-runtime-error")
-
-    assert res.status_code == 502
-    assert res.get_json() == {
-        "code": "plugin_runtime_error",
-        "message": "Plugin runtime request failed: Runtime.ExitError: Runtime exited with error: exit status 1",
-        "details": {
-            "request_id": "api-request-id",
-            "lambda_request_id": "lambda-request-id",
-        },
-        "status": 502,
-    }
-
-
 def test_external_api_param_mapping_and_quota():
     app = _create_api_app()
     client = app.test_client()

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

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

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

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

api/tests/unit_tests/services/test_agent_tool_inner_service.py

@@ -7,7 +7,6 @@ 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,
@@ -80,74 +79,6 @@ 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()
@@ -180,74 +111,6 @@ def test_invoke_uses_agent_tool_runtime_and_returns_observation(sqlite_session:
     assert sqlite_session.in_transaction()
 
 
-@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
-def test_invoke_keeps_ui_in_messages_and_out_of_observation(sqlite_session: Session) -> None:
-    _persist_app(sqlite_session)
-
-    with (
-        patch("services.agent_tool_inner_service.ToolManager.get_agent_tool_runtime", return_value=MagicMock()),
-        patch("services.agent_tool_inner_service.ToolEngine.generic_invoke", return_value=_ui_messages()),
-        patch(
-            "services.agent_tool_inner_service.ToolFileMessageTransformer.transform_tool_invoke_messages",
-            side_effect=lambda messages, **_kwargs: messages,
-        ),
-    ):
-        response = AgentToolInnerService().invoke(_request(), session=sqlite_session)
-
-    assert response.observation == "Sunny"
-    assert response.messages[1] == {
-        "type": "ui",
-        "message": {
-            "protocol": "a2ui",
-            "protocol_version": "v0.9.1",
-            "messages": [
-                {
-                    "version": "v0.9.1",
-                    "createSurface": {
-                        "surfaceId": "weather",
-                        "catalogId": "https://dify.ai/a2ui/catalog/v1",
-                    },
-                },
-                {
-                    "version": "v0.9.1",
-                    "updateComponents": {
-                        "surfaceId": "weather",
-                        "components": [
-                            {
-                                "id": "root",
-                                "component": "Text",
-                                "text": "Sunny",
-                            }
-                        ],
-                    },
-                },
-            ],
-        },
-        "meta": None,
-    }
-
-
-@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
-def test_invoke_drops_oversized_ui_batch_without_losing_observation(sqlite_session: Session) -> None:
-    _persist_app(sqlite_session)
-
-    with (
-        patch("services.agent_tool_inner_service.ToolManager.get_agent_tool_runtime", return_value=MagicMock()),
-        patch(
-            "services.agent_tool_inner_service.ToolEngine.generic_invoke",
-            return_value=_oversized_ui_messages(),
-        ),
-        patch(
-            "services.agent_tool_inner_service.ToolFileMessageTransformer.transform_tool_invoke_messages",
-            side_effect=lambda messages, **_kwargs: messages,
-        ),
-    ):
-        response = AgentToolInnerService().invoke(_request(), session=sqlite_session)
-
-    assert response.observation == "Sunny"
-    assert [message["type"] for message in response.messages] == ["text"]
-
-
 @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True)
 def test_invoke_raises_app_not_found_when_session_has_no_app(sqlite_session: Session) -> None:
     with pytest.raises(AgentToolInnerServiceError) as exc_info:

api/uv.lock

@@ -1371,6 +1371,7 @@ dependencies = [
     { name = "python-socketio" },
     { name = "readabilipy" },
     { name = "redis", extra = ["hiredis"] },
+    { name = "redis-entraid" },
     { name = "resend" },
     { name = "sendgrid" },
     { name = "sseclient-py" },
@@ -1655,6 +1656,7 @@ requires-dist = [
     { name = "python-socketio", specifier = ">=5.13.0,<6.0.0" },
     { name = "readabilipy", specifier = "==0.3.0" },
     { name = "redis", extras = ["hiredis"], specifier = ">=7.4.0,<8.0.0" },
+    { name = "redis-entraid", specifier = ">=1.2.0,<2.0.0" },
     { name = "resend", specifier = ">=2.27.0,<3.0.0" },
     { name = "sendgrid", specifier = ">=6.12.5,<7.0.0" },
     { name = "sseclient-py", specifier = ">=1.8.0,<2.0.0" },
@@ -5819,6 +5821,22 @@ hiredis = [
     { name = "hiredis" },
 ]
 
+[[package]]
+name = "redis-entraid"
+version = "1.2.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "azure-identity" },
+    { name = "msal" },
+    { name = "pyjwt" },
+    { name = "redis" },
+    { name = "requests" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a2/a7/0ddaeb27b33c76709e05a12b3bbeefce893c82a3a830146608d6fe620000/redis_entraid-1.2.1.tar.gz", hash = "sha256:a7c479ce46e6edb35bce9dd804d1cad7be99a3330815cfe028a648b486a10b41", size = 9792, upload-time = "2026-06-03T11:38:55.613Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cd/ca/01b8607102de756b270d3f6befeee700bd82dace4303d6f47ce0f53c11b0/redis_entraid-1.2.1-py3-none-any.whl", hash = "sha256:9de7e4a716b156d966a2d6bb5b5ccd64a692db30ae21fe3987f57d233793d558", size = 7967, upload-time = "2026-06-03T11:38:54.497Z" },
+]
+
 [[package]]
 name = "referencing"
 version = "0.37.0"

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

@@ -9,14 +9,12 @@ 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, ToolReturn
+from pydantic_ai import RunContext, Tool
 from pydantic_ai.tools import ToolDefinition
 from typing_extensions import Self, override
 
@@ -36,10 +34,6 @@ 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):
@@ -110,19 +104,13 @@ 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 | ToolReturn:
+        async def invoke_tool(_ctx: RunContext[object], **tool_arguments: object) -> str:
             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."
@@ -153,86 +141,6 @@ 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,14 +17,12 @@ 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, ConfigDict, Field, ValidationInfo, field_validator, model_validator
+from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator
 
 from dify_agent.layers.dify_plugin.configs import DifyPluginToolCredentialType
 from dify_agent.plugin_daemon_transport import (
@@ -64,32 +62,6 @@ 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
 
@@ -136,19 +108,10 @@ class DifyPluginToolInvokeMessage(BaseModel):
         FILE = "file"
         LOG = "log"
         BLOB_CHUNK = "blob_chunk"
-        UI = "ui"
 
     type: MessageType = MessageType.TEXT
     message: (
-        TextMessage
-        | JsonMessage
-        | UIMessage
-        | BlobChunkMessage
-        | BlobMessage
-        | LogMessage
-        | FileMessage
-        | VariableMessage
-        | None
+        TextMessage | JsonMessage | 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, ToolReturn
+from pydantic_ai import RunContext, Tool
 from pydantic_ai.tools import ToolDefinition
 from typing_extensions import Self, override
 
@@ -56,9 +56,6 @@ _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):
@@ -262,7 +259,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 | ToolReturn:
+    async def invoke_tool(_ctx: RunContext[object], **tool_arguments: object) -> str:
         try:
             merged_arguments = await _prepare_tool_arguments(
                 effective_parameters,
@@ -277,13 +274,7 @@ def _build_pydantic_ai_tool(
                 credentials=dict(tool_config.credentials),
                 tool_parameters=merged_arguments,
             )
-            observation, ui_messages = _convert_tool_response(messages)
-            if ui_messages:
-                return ToolReturn(
-                    observation,
-                    metadata={"dify_ui_messages": ui_messages},
-                )
-            return observation
+            return _convert_tool_response_to_text(messages)
         except DifyPluginToolClientError as exc:
             return _tool_error_text(tool_name=tool_name, error=exc)
         except ValueError as exc:
@@ -652,56 +643,8 @@ 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:
@@ -722,22 +665,9 @@ def _convert_tool_response(
             )
         elif response.type is DifyPluginToolInvokeMessage.MessageType.JSON:
             json_message = response.message
-            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)
+            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))
         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))
@@ -745,7 +675,7 @@ def _convert_tool_response(
     if json_parts:
         existing_parts = set(parts)
         parts.extend(part for part in json_parts if part not in existing_parts)
-    return "".join(parts), [] if invalid_ui_batch else ui_messages
+    return "".join(parts)
 
 
 __all__ = ["DifyPluginToolsDeps", "DifyPluginToolsLayer"]

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

@@ -1,11 +1,9 @@
 import asyncio
-import json
 import sys
 import types
 
 import httpx
 import pytest
-from pydantic_ai import ToolReturn
 
 from agenton.compositor import Compositor, LayerNode, LayerProvider
 
@@ -54,7 +52,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, _extract_ui_messages  # noqa: E402
+from dify_agent.layers.dify_core_tools.layer import DifyCoreToolsLayer  # noqa: E402
 from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig  # noqa: E402
 from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer  # noqa: E402
 
@@ -99,54 +97,6 @@ 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(
@@ -193,117 +143,6 @@ 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,12 +1,10 @@
 import asyncio
 import json
-import math
 from types import SimpleNamespace
 
 import httpx
 import pytest
-from pydantic import JsonValue, ValidationError
-from pydantic_ai import ToolReturn
+from pydantic import JsonValue
 
 from agenton.compositor import Compositor, LayerNode, LayerProvider
 from dify_agent.adapters.llm import DifyLLMAdapterModel
@@ -22,12 +20,7 @@ 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.tool_client import DifyPluginToolInvokeMessage
-from dify_agent.layers.dify_plugin.tools_layer import (
-    DifyPluginToolsLayer,
-    _convert_tool_response,
-    _PluginToolFileContext,
-)
+from dify_agent.layers.dify_plugin.tools_layer import DifyPluginToolsLayer, _PluginToolFileContext
 from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
 from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer
 
@@ -210,7 +203,6 @@ 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)
@@ -225,43 +217,6 @@ 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 '}}})}",
@@ -276,7 +231,6 @@ 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"):
@@ -291,68 +245,13 @@ def _tool_transport(
                 "api_version": "2026-01",
                 "auth_scope": "workspace",
             }
-            return _invoke_stream_response(
-                error_payload=invoke_error_payload,
-                chunked_blob=chunked_blob,
-                ui_message=ui_message,
-            )
+            return _invoke_stream_response(error_payload=invoke_error_payload, chunked_blob=chunked_blob)
 
         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,
@@ -442,188 +341,6 @@ def test_dify_plugin_tools_layer_uses_prepared_tool_definition_and_invokes_daemo
     asyncio.run(scenario())
 
 
-def test_dify_plugin_tools_layer_keeps_ui_out_of_model_observation() -> None:
-    async def scenario() -> None:
-        compositor = Compositor(
-            [
-                LayerNode("execution_context", _execution_context_provider()),
-                LayerNode("tools", _tools_provider(), deps={"execution_context": "execution_context"}),
-            ]
-        )
-        async with httpx.AsyncClient(transport=_tool_transport(ui_message=True)) as client:
-            async with compositor.enter(
-                configs={"execution_context": _execution_context_config(), "tools": _tools_config()}
-            ) as run:
-                tool = (
-                    await run.get_layer("tools", DifyPluginToolsLayer).get_tools(
-                        http_client=client,
-                        dify_api_http_client=client,
-                    )
-                )[0]
-                result = await tool.function_schema.call(
-                    {"query": "dify", "region": "global"},
-                    None,  # pyright: ignore[reportArgumentType]
-                )
-
-        assert isinstance(result, ToolReturn)
-        assert result.return_value == "Sunny"
-        assert result.metadata == {
-            "dify_ui_messages": [
-                {
-                    "protocol": "a2ui",
-                    "protocol_version": "v0.9.1",
-                    "messages": [
-                        {
-                            "version": "v0.9.1",
-                            "createSurface": {
-                                "surfaceId": "weather",
-                                "catalogId": "https://dify.ai/a2ui/catalog/v1",
-                            },
-                        },
-                        {
-                            "version": "v0.9.1",
-                            "updateComponents": {
-                                "surfaceId": "weather",
-                                "components": [
-                                    {
-                                        "id": "root",
-                                        "component": "Text",
-                                        "text": "Sunny",
-                                    }
-                                ],
-                            },
-                        },
-                    ],
-                    "fallback": "Sunny",
-                }
-            ]
-        }
-
-    asyncio.run(scenario())
-
-
-def test_dify_plugin_tools_layer_supports_safe_variable_ui_fallback() -> None:
-    ui_message = _ui_payload("weather")
-    observation, ui_messages = _convert_tool_response(
-        [
-            DifyPluginToolInvokeMessage(
-                type=DifyPluginToolInvokeMessage.MessageType.VARIABLE,
-                message=DifyPluginToolInvokeMessage.VariableMessage(
-                    variable_name="__dify_ui__",
-                    variable_value=ui_message,
-                ),
-            )
-        ]
-    )
-
-    assert observation == ""
-    assert ui_messages == [ui_message]
-
-
-def test_dify_plugin_ui_message_payload_size_includes_full_envelope() -> None:
-    with pytest.raises(ValidationError, match="128 KiB"):
-        DifyPluginToolInvokeMessage.UIMessage(
-            protocol="a2ui",
-            protocol_version="v0.9.1",
-            messages=[{"value": "x" * 127_000}],
-            fallback="y" * 4096,
-        )
-
-
-def test_dify_plugin_ui_message_rejects_non_json_values() -> None:
-    with pytest.raises(ValidationError, match="JSON-serializable"):
-        DifyPluginToolInvokeMessage.UIMessage(
-            protocol="a2ui",
-            protocol_version="v0.9.1",
-            messages=[{"value": object()}],
-        )
-
-
-@pytest.mark.parametrize(
-    "value",
-    [float("nan"), float("inf"), float("-inf")],
-    ids=["nan", "positive-infinity", "negative-infinity"],
-)
-def test_dify_plugin_ui_message_rejects_non_finite_numbers_without_restricting_json_messages(
-    value: float,
-) -> None:
-    with pytest.raises(ValidationError, match="finite"):
-        DifyPluginToolInvokeMessage.UIMessage(
-            protocol="a2ui",
-            protocol_version="v0.9.1",
-            messages=[{"value": value}],
-        )
-
-    json_message = DifyPluginToolInvokeMessage.JsonMessage(json_object={"value": value})
-    assert isinstance(json_message.json_object, dict)
-    stored_value = json_message.json_object["value"]
-    assert isinstance(stored_value, float)
-    assert not math.isfinite(stored_value)
-
-
-def test_dify_plugin_tools_layer_drops_ui_batch_on_seventeenth_message() -> None:
-    observation, ui_messages = _convert_tool_response(
-        [
-            DifyPluginToolInvokeMessage(
-                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
-                message=DifyPluginToolInvokeMessage.TextMessage(text="before"),
-            ),
-            *[_native_ui_message(_ui_payload(f"surface-{index}")) for index in range(17)],
-            DifyPluginToolInvokeMessage(
-                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
-                message=DifyPluginToolInvokeMessage.TextMessage(text=" after"),
-            ),
-        ]
-    )
-
-    assert observation == "before after"
-    assert ui_messages == []
-
-
-def test_dify_plugin_tools_layer_drops_ui_batch_on_aggregate_payload_overflow() -> None:
-    payloads = [_large_ui_payload(f"surface-{index}") for index in range(6)]
-    assert len(json.dumps(payloads, ensure_ascii=False, separators=(",", ":")).encode()) > 512 * 1024
-
-    observation, ui_messages = _convert_tool_response(
-        [
-            DifyPluginToolInvokeMessage(
-                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
-                message=DifyPluginToolInvokeMessage.TextMessage(text="weather observation"),
-            ),
-            *[_native_ui_message(payload) for payload in payloads],
-        ]
-    )
-
-    assert observation == "weather observation"
-    assert ui_messages == []
-
-
-@pytest.mark.parametrize(
-    "invalid_payload",
-    ["invalid", {"protocol": "invalid"}],
-    ids=["non-object", "invalid-object"],
-)
-def test_dify_plugin_tools_layer_suppresses_invalid_reserved_ui_envelopes(
-    invalid_payload: object,
-) -> None:
-    observation, ui_messages = _convert_tool_response(
-        [
-            DifyPluginToolInvokeMessage(
-                type=DifyPluginToolInvokeMessage.MessageType.TEXT,
-                message=DifyPluginToolInvokeMessage.TextMessage(text="weather observation"),
-            ),
-            _native_ui_message(_ui_payload("valid")),
-            DifyPluginToolInvokeMessage(
-                type=DifyPluginToolInvokeMessage.MessageType.JSON,
-                message=DifyPluginToolInvokeMessage.JsonMessage(json_object={"__dify_ui__": invalid_payload}),
-            ),
-        ]
-    )
-
-    assert observation == "weather observation"
-    assert ui_messages == []
-
-
 def test_dify_plugin_tools_layer_uses_each_tool_plugin_id_for_transport() -> None:
     async def scenario() -> None:
         seen_requests: list[tuple[str, str, str, str]] = []

docker/.env.example

@@ -121,6 +121,7 @@ REDIS_KEEPALIVE=true
 REDIS_KEEPALIVE_IDLE=30
 REDIS_KEEPALIVE_INTERVAL=10
 REDIS_KEEPALIVE_COUNT=10
+REDIS_USE_AZURE_MANAGED_IDENTITY=false
 CELERY_BROKER_URL=redis://:difyai123456@redis:6379/1
 CELERY_BACKEND=redis
 BROKER_USE_SSL=false

oxlint-suppressions.json

@@ -835,6 +835,14 @@
       "count": 1
     }
   },
+  "web/app/components/base/audio-btn/audio.ts": {
+    "node-js/prefer-global/buffer": {
+      "count": 1
+    },
+    "typescript/no-explicit-any": {
+      "count": 3
+    }
+  },
   "web/app/components/base/audio-gallery/AudioPlayer.tsx": {
     "jsx_a11y/media-has-caption": {
       "count": 1

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -1,5 +0,0 @@
-import { KnowledgeRoutePlaceholder } from '@/features/new-rag/knowledge-route-placeholder'
-
-export default function Page() {
-  return <KnowledgeRoutePlaceholder type="documents" />
-}

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

@@ -1,18 +0,0 @@
-import { KnowledgeRouteGuard } from '@/features/new-rag/knowledge-route-guard'
-import { KnowledgeSpaceShell } from '@/features/new-rag/knowledge-space-shell'
-
-export default async function Layout({
-  children,
-  params,
-}: {
-  children: React.ReactNode
-  params: Promise<{ knowledgeSpaceId: string }>
-}) {
-  const { knowledgeSpaceId } = await params
-
-  return (
-    <KnowledgeRouteGuard>
-      <KnowledgeSpaceShell knowledgeSpaceId={knowledgeSpaceId}>{children}</KnowledgeSpaceShell>
-    </KnowledgeRouteGuard>
-  )
-}

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

@@ -1,6 +0,0 @@
-import { redirect } from '@/next/navigation'
-
-export default async function Page({ params }: { params: Promise<{ knowledgeSpaceId: string }> }) {
-  const { knowledgeSpaceId } = await params
-  redirect(`/datasets/new/${knowledgeSpaceId}/sources`)
-}

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

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

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

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

web/app/(commonLayout)/datasets/new/create/page.tsx

@@ -1,10 +0,0 @@
-import { CreateKnowledgePage } from '@/features/new-rag/create-knowledge-page'
-import { KnowledgeRouteGuard } from '@/features/new-rag/knowledge-route-guard'
-
-export default function Page() {
-  return (
-    <KnowledgeRouteGuard>
-      <CreateKnowledgePage />
-    </KnowledgeRouteGuard>
-  )
-}

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

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

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

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

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

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

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

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