dify/api/services/workflow_runtime/chatflow_pipeline.py

import json
import logging
import time
from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from threading import Thread
from typing import Any

from sqlalchemy import select, update
from sqlalchemy.orm import Session

from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
from core.app.apps.common.graph_runtime_state_support import GraphRuntimeStateSupport
from core.app.apps.draft_variable_saver import DraftVariableSaverFactory
from core.app.entities.app_invoke_entities import (
    AdvancedChatAppGenerateEntity,
    InvokeFrom,
    get_credit_usage_app_type,
)
from core.app.entities.queue_entities import (
    MessageQueueMessage,
    QueueAdvancedChatMessageEndEvent,
    QueueAgentLogEvent,
    QueueAnnotationReplyEvent,
    QueueErrorEvent,
    QueueHumanInputFormFilledEvent,
    QueueHumanInputFormTimeoutEvent,
    QueueIterationCompletedEvent,
    QueueIterationNextEvent,
    QueueIterationStartEvent,
    QueueLoopCompletedEvent,
    QueueLoopNextEvent,
    QueueLoopStartEvent,
    QueueMessageReplaceEvent,
    QueueNodeExceptionEvent,
    QueueNodeFailedEvent,
    QueueNodeRetryEvent,
    QueueNodeStartedEvent,
    QueueNodeSucceededEvent,
    QueuePingEvent,
    QueueReasoningChunkEvent,
    QueueRetrieverResourcesEvent,
    QueueStopEvent,
    QueueTextChunkEvent,
    QueueWorkflowFailedEvent,
    QueueWorkflowPartialSuccessEvent,
    QueueWorkflowPausedEvent,
    QueueWorkflowStartedEvent,
    QueueWorkflowSucceededEvent,
    WorkflowQueueMessage,
)
from core.app.entities.task_entities import (
    AdvancedChatPausedBlockingResponse,
    ChatbotAppBlockingResponse,
    ChatbotAppStreamResponse,
    ErrorStreamResponse,
    HumanInputRequiredPauseReasonPayload,
    HumanInputRequiredResponse,
    MessageAudioEndStreamResponse,
    MessageAudioStreamResponse,
    MessageEndStreamResponse,
    PingStreamResponse,
    ReasoningChunkStreamResponse,
    StreamResponse,
    WorkflowPauseStreamResponse,
    WorkflowTaskState,
)
from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
from core.app.task_pipeline.message_cycle_manager import MessageCycleManager
from core.base.tts import AppGeneratorTTSPublisher
from core.db.session_factory import session_factory
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.workflow.file_reference import resolve_file_record_id
from core.workflow.nodes.human_input.pause_reason import HumanInputRequired
from core.workflow.system_variables import build_system_variables
from graphon.enums import WorkflowExecutionStatus
from graphon.model_runtime.entities.llm_entities import LLMUsage
from graphon.model_runtime.utils.encoders import jsonable_encoder
from graphon.nodes import BuiltinNodeTypes
from graphon.runtime import RuntimeState
from libs.datetime_utils import naive_utc_now
from models import Account, Conversation, EndUser, Message, MessageFile
from models.enums import CreatorUserRole, MessageFileBelongsTo, MessageStatus
from models.execution_extra_content import HumanInputContent
from models.model import AppMode
from models.workflow import Workflow, WorkflowRun
from services.workflow_runtime.response_converter import WorkflowResponseConverter

logger = logging.getLogger(__name__)


@dataclass(frozen=True, slots=True)
class WorkflowSnapshot:
    id: str
    tenant_id: str
    features_dict: Mapping[str, Any]

    @classmethod
    def from_workflow(cls, workflow: Workflow) -> "WorkflowSnapshot":
        return cls(
            id=workflow.id,
            tenant_id=workflow.tenant_id,
            features_dict=dict(workflow.features_dict),
        )


@dataclass(frozen=True, slots=True)
class ConversationSnapshot:
    id: str
    mode: AppMode

    @classmethod
    def from_conversation(cls, conversation: Conversation) -> "ConversationSnapshot":
        return cls(
            id=conversation.id,
            mode=conversation.mode,
        )


@dataclass(frozen=True, slots=True)
class MessageSnapshot:
    id: str
    query: str
    created_at: datetime
    status: MessageStatus
    answer: str

    @classmethod
    def from_message(cls, message: Message) -> "MessageSnapshot":
        return cls(
            id=message.id,
            query=message.query,
            created_at=message.created_at,
            status=message.status,
            answer=message.answer,
        )


class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
    """
    AdvancedChatAppGenerateTaskPipeline is a class that generate stream output and state management for Application.
    """

    def __init__(
        self,
        application_generate_entity: AdvancedChatAppGenerateEntity,
        workflow: WorkflowSnapshot,
        queue_manager: AppQueueManager,
        conversation: ConversationSnapshot,
        message: MessageSnapshot,
        user: Account | EndUser,
        stream: bool,
        dialogue_count: int,
        draft_var_saver_factory: DraftVariableSaverFactory,
    ):
        self._base_task_pipeline = BasedGenerateTaskPipeline(
            application_generate_entity=application_generate_entity,
            queue_manager=queue_manager,
            stream=stream,
        )

        match user:
            case EndUser():
                self._user_id = user.id
                user_session_id = user.session_id
                self._created_by_role = CreatorUserRole.END_USER
            case Account():
                self._user_id = user.id
                user_session_id = user.id
                self._created_by_role = CreatorUserRole.ACCOUNT
            case _:
                raise NotImplementedError(f"User type not supported: {type(user)}")

        self._workflow_system_variables = build_system_variables(
            query=message.query,
            files=application_generate_entity.files,
            conversation_id=conversation.id,
            user_id=user_session_id,
            dialogue_count=dialogue_count,
            app_id=application_generate_entity.app_config.app_id,
            workflow_id=workflow.id,
            workflow_execution_id=application_generate_entity.workflow_run_id,
        )
        self._workflow_response_converter = WorkflowResponseConverter(
            application_generate_entity=application_generate_entity,
            user=user,
            system_variables=self._workflow_system_variables,
        )

        self._task_state = WorkflowTaskState()
        self._seed_task_state_from_message(message)
        self._message_cycle_manager = MessageCycleManager(
            application_generate_entity=application_generate_entity, task_state=self._task_state
        )

        self._application_generate_entity = application_generate_entity
        self._workflow_id = workflow.id
        self._workflow_features_dict = workflow.features_dict
        self._workflow_tenant_id = workflow.tenant_id
        self._conversation_id = conversation.id
        self._conversation_mode = conversation.mode
        self._message_id = message.id
        self._message_created_at = int(message.created_at.timestamp())
        self._conversation_name_generate_thread: Thread | None = None
        self._recorded_files: list[Mapping[str, Any]] = []
        self._workflow_run_id: str = ""
        self._draft_var_saver_factory = draft_var_saver_factory
        self._graph_runtime_state: RuntimeState | None = None
        self._message_saved_on_pause = False
        self._seed_graph_runtime_state_from_queue_manager()

    def _seed_task_state_from_message(self, message: MessageSnapshot) -> None:
        if message.status == MessageStatus.PAUSED and message.answer:
            self._task_state.answer = message.answer

    def process(
        self,
    ) -> ChatbotAppBlockingResponse | AdvancedChatPausedBlockingResponse | Generator[ChatbotAppStreamResponse]:
        """
        Process generate task pipeline.
        :return:
        """
        self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name(
            conversation_id=self._conversation_id,
            query=self._application_generate_entity.query,
            message_id=self._message_id,
        )

        generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager)

        if self._base_task_pipeline.stream:
            return self._to_stream_response(generator)
        else:
            return self._to_blocking_response(generator)

    def _to_blocking_response(
        self, generator: Generator[StreamResponse]
    ) -> ChatbotAppBlockingResponse | AdvancedChatPausedBlockingResponse:
        """
        Process blocking response.
        :return:
        """
        human_input_responses: list[HumanInputRequiredResponse] = []
        for stream_response in generator:
            match stream_response:
                case ErrorStreamResponse():
                    raise stream_response.err
                case HumanInputRequiredResponse():
                    human_input_responses.append(stream_response)
                case WorkflowPauseStreamResponse():
                    return AdvancedChatPausedBlockingResponse(
                        task_id=stream_response.task_id,
                        data=AdvancedChatPausedBlockingResponse.Data(
                            id=self._message_id,
                            mode=self._conversation_mode,
                            conversation_id=self._conversation_id,
                            message_id=self._message_id,
                            workflow_run_id=stream_response.data.workflow_run_id,
                            answer=self._task_state.answer,
                            metadata=self._message_end_to_stream_response().metadata,
                            created_at=self._message_created_at,
                            paused_nodes=stream_response.data.paused_nodes,
                            reasons=stream_response.data.reasons,
                            status=stream_response.data.status,
                            elapsed_time=stream_response.data.elapsed_time,
                            total_tokens=stream_response.data.total_tokens,
                            total_steps=stream_response.data.total_steps,
                        ),
                    )
                case MessageEndStreamResponse():
                    extras = {}
                    if stream_response.metadata:
                        extras["metadata"] = stream_response.metadata

                    return ChatbotAppBlockingResponse(
                        task_id=stream_response.task_id,
                        data=ChatbotAppBlockingResponse.Data(
                            id=self._message_id,
                            mode=self._conversation_mode,
                            conversation_id=self._conversation_id,
                            message_id=self._message_id,
                            answer=self._task_state.answer,
                            created_at=self._message_created_at,
                            **extras,
                        ),
                    )
                case _:
                    continue

        if human_input_responses:
            return self._build_paused_blocking_response_from_human_input(human_input_responses)

        raise ValueError("queue listening stopped unexpectedly.")

    def _build_paused_blocking_response_from_human_input(
        self, human_input_responses: list[HumanInputRequiredResponse]
    ) -> AdvancedChatPausedBlockingResponse:
        runtime_state = self._resolve_graph_runtime_state()
        paused_nodes = list(dict.fromkeys(response.data.node_id for response in human_input_responses))
        reasons = [
            HumanInputRequiredPauseReasonPayload.from_response_data(response.data).model_dump(mode="json")
            for response in human_input_responses
        ]

        return AdvancedChatPausedBlockingResponse(
            task_id=self._application_generate_entity.task_id,
            data=AdvancedChatPausedBlockingResponse.Data(
                id=self._message_id,
                mode=self._conversation_mode,
                conversation_id=self._conversation_id,
                message_id=self._message_id,
                workflow_run_id=human_input_responses[-1].workflow_run_id,
                answer=self._task_state.answer,
                metadata=self._message_end_to_stream_response().metadata,
                created_at=self._message_created_at,
                paused_nodes=paused_nodes,
                reasons=reasons,
                status=WorkflowExecutionStatus.PAUSED,
                elapsed_time=self._application_generate_entity.workflow_execution_timing.finish(),
                total_tokens=runtime_state.total_tokens,
                total_steps=runtime_state.node_run_steps,
            ),
        )

    def _to_stream_response(self, generator: Generator[StreamResponse]) -> Generator[ChatbotAppStreamResponse, Any]:
        """
        To stream response.
        :return:
        """
        for stream_response in generator:
            yield ChatbotAppStreamResponse(
                conversation_id=self._conversation_id,
                message_id=self._message_id,
                created_at=self._message_created_at,
                stream_response=stream_response,
            )

    def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: str):
        if not publisher:
            return None
        audio_msg = publisher.check_and_get_audio()
        if audio_msg is None:
            return None
        if audio_msg.status == "responding":
            return MessageAudioStreamResponse(audio=audio_msg.audio, audio_type=audio_msg.audio_type, task_id=task_id)
        if audio_msg.status in {"finish", "error"}:
            return None
        raise RuntimeError(f"TTS publisher returned an unknown status: {audio_msg.status}")

    def _wrapper_process_stream_response(
        self, trace_manager: TraceQueueManager | None = None
    ) -> Generator[StreamResponse]:
        tts_publisher = None
        task_id = self._application_generate_entity.task_id
        tenant_id = self._application_generate_entity.app_config.tenant_id
        features_dict = self._workflow_features_dict

        if (
            self._base_task_pipeline.stream
            and features_dict.get("text_to_speech")
            and features_dict["text_to_speech"].get("enabled")
            and features_dict["text_to_speech"].get("autoPlay") == "enabled"
        ):
            tts_publisher = AppGeneratorTTSPublisher(
                tenant_id,
                features_dict["text_to_speech"].get("voice"),
                features_dict["text_to_speech"].get("language"),
                get_credit_usage_app_type(self._application_generate_entity.app_config.app_mode),
            )

        try:
            for response in self._process_stream_response(tts_publisher=tts_publisher, trace_manager=trace_manager):
                while audio_response := self._listen_audio_msg(publisher=tts_publisher, task_id=task_id):
                    yield audio_response
                if tts_publisher and isinstance(response, ErrorStreamResponse):
                    tts_publisher.cancel()
                    yield MessageAudioEndStreamResponse(audio="", task_id=task_id)
                    yield response
                    return
                yield response

            if tts_publisher is None:
                return

            tts_publisher.publish(None)
            while True:
                audio_trunk = tts_publisher.check_and_get_audio(block=True)
                assert audio_trunk is not None
                if audio_trunk.status == "responding":
                    yield MessageAudioStreamResponse(
                        audio=audio_trunk.audio, audio_type=audio_trunk.audio_type, task_id=task_id
                    )
                    continue
                if audio_trunk.status not in {"finish", "error"}:
                    raise RuntimeError(f"TTS publisher returned an unknown status: {audio_trunk.status}")

                yield MessageAudioEndStreamResponse(audio="", task_id=task_id)
                if audio_trunk.status == "error":
                    if audio_trunk.error is None:
                        raise RuntimeError("TTS publisher returned an error terminal without an exception")
                    yield ErrorStreamResponse(err=audio_trunk.error, task_id=task_id)
                return
        finally:
            if tts_publisher:
                tts_publisher.cancel()

    @contextmanager
    def _database_session(self):
        """Context manager for database sessions."""
        with session_factory.create_session() as session:
            try:
                yield session
                session.commit()
            except Exception:
                session.rollback()
                raise

    def _ensure_workflow_initialized(self):
        """Fluent validation for workflow state."""
        if not self._workflow_run_id:
            raise ValueError("workflow run not initialized.")

    def _handle_ping_event(self, event: QueuePingEvent, **kwargs) -> Generator[PingStreamResponse]:
        """Handle ping events."""
        yield self._base_task_pipeline.ping_stream_response()

    def _handle_error_event(self, event: QueueErrorEvent, **kwargs) -> Generator[ErrorStreamResponse]:
        """Handle error events."""
        with self._database_session() as session:
            err = self._base_task_pipeline.handle_error(event=event, session=session, message_id=self._message_id)
        yield self._base_task_pipeline.error_to_stream_response(err)

    def _handle_workflow_started_event(
        self,
        event: QueueWorkflowStartedEvent,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle workflow started events."""
        runtime_state = self._resolve_graph_runtime_state()
        run_id = self._extract_workflow_run_id(runtime_state)
        self._workflow_run_id = run_id

        with self._database_session() as session:
            session.execute(update(Message).where(Message.id == self._message_id).values(workflow_run_id=run_id))

        workflow_start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_run_id=run_id,
            workflow_id=self._workflow_id,
            reason=event.reason,
            node_execution_snapshots=event.node_execution_snapshots,
        )

        yield workflow_start_resp

    def _handle_node_retry_event(self, event: QueueNodeRetryEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle node retry events."""
        self._ensure_workflow_initialized()

        node_retry_resp = self._workflow_response_converter.workflow_node_retry_to_stream_response(
            event=event,
            task_id=self._application_generate_entity.task_id,
        )

        if node_retry_resp:
            yield node_retry_resp

    def _handle_node_started_event(self, event: QueueNodeStartedEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle node started events."""
        self._ensure_workflow_initialized()

        node_start_resp = self._workflow_response_converter.workflow_node_start_to_stream_response(
            event=event,
            task_id=self._application_generate_entity.task_id,
        )

        if node_start_resp:
            yield node_start_resp

    def _handle_node_succeeded_event(self, event: QueueNodeSucceededEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle node succeeded events."""
        # Record files if it's an answer node or end node
        if event.node_type in [BuiltinNodeTypes.ANSWER, BuiltinNodeTypes.END, BuiltinNodeTypes.LLM]:
            self._recorded_files.extend(
                self._workflow_response_converter.fetch_files_from_node_outputs(event.outputs or {})
            )

        # Collect terminal reasoning (separated mode) per LLM node id for persistence. This is the
        # authoritative source (outputs.reasoning_content), decoupled from the live delta stream.
        # Accumulate across iteration/loop passes (same node_id) to match the live stream, which
        # appends every pass under the same key — overwriting would keep only the last pass.
        if event.node_type == BuiltinNodeTypes.LLM:
            reasoning_content = (event.outputs or {}).get("reasoning_content")
            if isinstance(reasoning_content, str) and reasoning_content:
                self._task_state.metadata.reasoning[event.node_id] = (
                    self._task_state.metadata.reasoning.get(event.node_id, "") + reasoning_content
                )

        node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response(
            event=event,
            task_id=self._application_generate_entity.task_id,
        )

        self._save_output_for_event(event, event.node_execution_id)

        if node_finish_resp:
            yield node_finish_resp

    def _handle_node_failed_events(
        self,
        event: QueueNodeFailedEvent | QueueNodeExceptionEvent,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle various node failure events."""
        node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response(
            event=event,
            task_id=self._application_generate_entity.task_id,
        )

        if isinstance(event, QueueNodeExceptionEvent):
            self._save_output_for_event(event, event.node_execution_id)

        if node_finish_resp:
            yield node_finish_resp

    def _handle_text_chunk_event(
        self,
        event: QueueTextChunkEvent,
        *,
        tts_publisher: AppGeneratorTTSPublisher | None = None,
        queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle text chunk events."""
        delta_text = event.text
        if delta_text is None:
            return

        # Handle output moderation chunk
        should_direct_answer = self._handle_output_moderation_chunk(delta_text)
        if should_direct_answer:
            return

        current_time = time.perf_counter()
        if self._task_state.first_token_time is None and delta_text.strip():
            self._task_state.first_token_time = current_time
            self._task_state.is_streaming_response = True

        if delta_text.strip():
            self._task_state.last_token_time = current_time

        # Only publish tts message at text chunk streaming
        if tts_publisher and queue_message:
            tts_publisher.publish(queue_message)

        self._task_state.answer += delta_text
        yield self._message_cycle_manager.message_to_stream_response(
            answer=delta_text, message_id=self._message_id, from_variable_selector=event.from_variable_selector
        )

    def _handle_reasoning_chunk_event(self, event: QueueReasoningChunkEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle out-of-band reasoning chunk events.

        Pure emit: reasoning is streamed on its own channel and never written to the
        answer. The terminal marker (is_final) may carry an empty reasoning string, in
        which case it is still forwarded as the "thinking finished" signal.
        """
        if not event.reasoning and not event.is_final:
            return
        yield ReasoningChunkStreamResponse(
            task_id=self._application_generate_entity.task_id,
            data=ReasoningChunkStreamResponse.Data(
                message_id=self._message_id,
                reasoning=event.reasoning,
                node_id=event.from_node_id,
                is_final=event.is_final,
            ),
        )

    def _handle_iteration_start_event(self, event: QueueIterationStartEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle iteration start events."""
        self._ensure_workflow_initialized()

        iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_execution_id=self._workflow_run_id,
            event=event,
        )
        yield iter_start_resp

    def _handle_iteration_next_event(self, event: QueueIterationNextEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle iteration next events."""
        self._ensure_workflow_initialized()

        iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_execution_id=self._workflow_run_id,
            event=event,
        )
        yield iter_next_resp

    def _handle_iteration_completed_event(
        self, event: QueueIterationCompletedEvent, **kwargs
    ) -> Generator[StreamResponse]:
        """Handle iteration completed events."""
        self._ensure_workflow_initialized()

        iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_execution_id=self._workflow_run_id,
            event=event,
        )
        yield iter_finish_resp

    def _handle_loop_start_event(self, event: QueueLoopStartEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle loop start events."""
        self._ensure_workflow_initialized()

        loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_execution_id=self._workflow_run_id,
            event=event,
        )
        yield loop_start_resp

    def _handle_loop_next_event(self, event: QueueLoopNextEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle loop next events."""
        self._ensure_workflow_initialized()

        loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_execution_id=self._workflow_run_id,
            event=event,
        )
        yield loop_next_resp

    def _handle_loop_completed_event(self, event: QueueLoopCompletedEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle loop completed events."""
        self._ensure_workflow_initialized()

        loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_execution_id=self._workflow_run_id,
            event=event,
        )
        yield loop_finish_resp

    def _handle_workflow_succeeded_event(
        self,
        event: QueueWorkflowSucceededEvent,
        *,
        trace_manager: TraceQueueManager | None = None,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle workflow succeeded events."""
        _ = trace_manager
        self._ensure_workflow_initialized()
        validated_state = self._ensure_graph_runtime_initialized()
        workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_id=self._workflow_id,
            status=WorkflowExecutionStatus.SUCCEEDED,
            graph_runtime_state=validated_state,
        )

        yield from self._handle_advanced_chat_message_end_event(
            QueueAdvancedChatMessageEndEvent(), graph_runtime_state=validated_state
        )
        yield workflow_finish_resp

    def _handle_workflow_partial_success_event(
        self,
        event: QueueWorkflowPartialSuccessEvent,
        *,
        trace_manager: TraceQueueManager | None = None,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle workflow partial success events."""
        _ = trace_manager
        self._ensure_workflow_initialized()
        validated_state = self._ensure_graph_runtime_initialized()
        workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_id=self._workflow_id,
            status=WorkflowExecutionStatus.PARTIAL_SUCCEEDED,
            graph_runtime_state=validated_state,
            exceptions_count=event.exceptions_count,
        )

        yield from self._handle_advanced_chat_message_end_event(
            QueueAdvancedChatMessageEndEvent(), graph_runtime_state=validated_state
        )
        yield workflow_finish_resp

    def _handle_workflow_paused_event(
        self,
        event: QueueWorkflowPausedEvent,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle workflow paused events."""
        validated_state = self._ensure_graph_runtime_initialized()
        responses = self._workflow_response_converter.workflow_pause_to_stream_response(
            event=event,
            task_id=self._application_generate_entity.task_id,
            graph_runtime_state=validated_state,
        )
        for reason in event.reasons:
            if isinstance(reason, HumanInputRequired):
                self._persist_human_input_extra_content(form_id=reason.form_id)
        yield from responses
        resolved_state: RuntimeState | None = None
        try:
            resolved_state = self._ensure_graph_runtime_initialized()
        except ValueError:
            resolved_state = None

        with self._database_session() as session:
            # Serialize the message's pause state with durable stop on the run.
            workflow_run_status = session.scalar(
                select(WorkflowRun.status)
                .where(
                    WorkflowRun.id == self._workflow_run_id,
                    WorkflowRun.tenant_id == self._workflow_tenant_id,
                    WorkflowRun.app_id == self._application_generate_entity.app_config.app_id,
                )
                .with_for_update()
            )
            self._save_message(session=session, graph_runtime_state=resolved_state)
            message = self._get_message(session=session)
            if message is not None and workflow_run_status == WorkflowExecutionStatus.PAUSED:
                message.status = MessageStatus.PAUSED
            self._message_saved_on_pause = True
        self._base_task_pipeline.queue_manager.publish(QueueAdvancedChatMessageEndEvent(), PublishFrom.TASK_PIPELINE)

    def _handle_workflow_failed_event(
        self,
        event: QueueWorkflowFailedEvent,
        *,
        trace_manager: TraceQueueManager | None = None,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle workflow failed events."""
        _ = trace_manager
        self._ensure_workflow_initialized()
        validated_state = self._ensure_graph_runtime_initialized()

        workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
            task_id=self._application_generate_entity.task_id,
            workflow_id=self._workflow_id,
            status=WorkflowExecutionStatus.FAILED,
            graph_runtime_state=validated_state,
            error=event.error,
            exceptions_count=event.exceptions_count,
        )

        with self._database_session() as session:
            err_event = QueueErrorEvent(error=ValueError(f"Run failed: {event.error}"))
            err = self._base_task_pipeline.handle_error(event=err_event, session=session, message_id=self._message_id)

        yield workflow_finish_resp
        yield self._base_task_pipeline.error_to_stream_response(err)

    def _handle_stop_event(
        self,
        event: QueueStopEvent,
        *,
        graph_runtime_state: RuntimeState | None = None,
        trace_manager: TraceQueueManager | None = None,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle stop events."""
        _ = trace_manager
        resolved_state = None
        if self._workflow_run_id:
            resolved_state = self._resolve_graph_runtime_state(graph_runtime_state)

        if self._workflow_run_id and resolved_state:
            workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
                task_id=self._application_generate_entity.task_id,
                workflow_id=self._workflow_id,
                status=WorkflowExecutionStatus.STOPPED,
                graph_runtime_state=resolved_state,
                error=event.get_stop_reason(),
            )

            with self._database_session() as session:
                # Save message
                self._save_message(session=session, graph_runtime_state=resolved_state)
            self._emit_message_trace()

            yield workflow_finish_resp
        elif event.stopped_by in (
            QueueStopEvent.StopBy.INPUT_MODERATION,
            QueueStopEvent.StopBy.ANNOTATION_REPLY,
        ):
            # When hitting input-moderation or annotation-reply, the workflow will not start
            with self._database_session() as session:
                # Save message
                self._save_message(session=session)
            self._emit_message_trace()

        yield self._message_end_to_stream_response()

    def _handle_advanced_chat_message_end_event(
        self,
        event: QueueAdvancedChatMessageEndEvent,
        *,
        graph_runtime_state: RuntimeState | None = None,
        **kwargs,
    ) -> Generator[StreamResponse]:
        """Handle advanced chat message end events."""
        resolved_state = self._ensure_graph_runtime_initialized(graph_runtime_state)

        output_moderation_answer = self._base_task_pipeline.handle_output_moderation_when_task_finished(
            self._task_state.answer
        )
        if output_moderation_answer:
            self._task_state.answer = output_moderation_answer
            yield self._message_cycle_manager.message_replace_to_stream_response(
                answer=output_moderation_answer,
                reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION,
            )

        # Save message unless it has already been persisted on pause.
        if not self._message_saved_on_pause:
            with self._database_session() as session:
                self._save_message(session=session, graph_runtime_state=resolved_state)
            self._emit_message_trace()

        yield self._message_end_to_stream_response()

    def _handle_retriever_resources_event(
        self, event: QueueRetrieverResourcesEvent, **kwargs
    ) -> Generator[StreamResponse]:
        """Handle retriever resources events."""
        self._message_cycle_manager.handle_retriever_resources(event)
        yield from ()

    def _handle_annotation_reply_event(self, event: QueueAnnotationReplyEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle annotation reply events."""
        with self._database_session() as session:
            self._message_cycle_manager.handle_annotation_reply(event, session)
        yield from ()

    def _handle_message_replace_event(self, event: QueueMessageReplaceEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle message replace events."""
        yield self._message_cycle_manager.message_replace_to_stream_response(answer=event.text, reason=event.reason)

    def _handle_human_input_form_filled_event(
        self, event: QueueHumanInputFormFilledEvent, **kwargs
    ) -> Generator[StreamResponse]:
        """Handle human input form filled events."""
        self._persist_human_input_extra_content(form_id=event.form_id)
        yield self._workflow_response_converter.human_input_form_filled_to_stream_response(
            event=event, task_id=self._application_generate_entity.task_id
        )

    def _handle_human_input_form_timeout_event(
        self, event: QueueHumanInputFormTimeoutEvent, **kwargs
    ) -> Generator[StreamResponse]:
        """Handle human input form timeout events."""
        yield self._workflow_response_converter.human_input_form_timeout_to_stream_response(
            event=event, task_id=self._application_generate_entity.task_id
        )

    def _persist_human_input_extra_content(self, *, form_id: str) -> None:
        if not self._workflow_run_id or not self._message_id:
            return

        with self._database_session() as session:
            exists_stmt = select(HumanInputContent).where(
                HumanInputContent.workflow_run_id == self._workflow_run_id,
                HumanInputContent.message_id == self._message_id,
                HumanInputContent.form_id == form_id,
            )
            if session.scalar(exists_stmt) is not None:
                return

            content = HumanInputContent(
                workflow_run_id=self._workflow_run_id,
                message_id=self._message_id,
                form_id=form_id,
            )
            session.add(content)

    def _handle_agent_log_event(self, event: QueueAgentLogEvent, **kwargs) -> Generator[StreamResponse]:
        """Handle agent log events."""
        yield self._workflow_response_converter.handle_agent_log(
            task_id=self._application_generate_entity.task_id, event=event
        )

    def _get_event_handlers(self) -> dict[type, Callable]:
        """Get mapping of event types to their handlers using fluent pattern."""
        return {
            # Basic events
            QueuePingEvent: self._handle_ping_event,
            QueueErrorEvent: self._handle_error_event,
            QueueTextChunkEvent: self._handle_text_chunk_event,
            QueueReasoningChunkEvent: self._handle_reasoning_chunk_event,
            # Workflow events
            QueueWorkflowStartedEvent: self._handle_workflow_started_event,
            QueueWorkflowSucceededEvent: self._handle_workflow_succeeded_event,
            QueueWorkflowPartialSuccessEvent: self._handle_workflow_partial_success_event,
            QueueWorkflowPausedEvent: self._handle_workflow_paused_event,
            QueueWorkflowFailedEvent: self._handle_workflow_failed_event,
            # Node events
            QueueNodeRetryEvent: self._handle_node_retry_event,
            QueueNodeStartedEvent: self._handle_node_started_event,
            QueueNodeSucceededEvent: self._handle_node_succeeded_event,
            # Iteration events
            QueueIterationStartEvent: self._handle_iteration_start_event,
            QueueIterationNextEvent: self._handle_iteration_next_event,
            QueueIterationCompletedEvent: self._handle_iteration_completed_event,
            # Loop events
            QueueLoopStartEvent: self._handle_loop_start_event,
            QueueLoopNextEvent: self._handle_loop_next_event,
            QueueLoopCompletedEvent: self._handle_loop_completed_event,
            # Control events
            QueueStopEvent: self._handle_stop_event,
            # Message events
            QueueRetrieverResourcesEvent: self._handle_retriever_resources_event,
            QueueAnnotationReplyEvent: self._handle_annotation_reply_event,
            QueueMessageReplaceEvent: self._handle_message_replace_event,
            QueueAdvancedChatMessageEndEvent: self._handle_advanced_chat_message_end_event,
            QueueAgentLogEvent: self._handle_agent_log_event,
            QueueHumanInputFormFilledEvent: self._handle_human_input_form_filled_event,
            QueueHumanInputFormTimeoutEvent: self._handle_human_input_form_timeout_event,
        }

    def _dispatch_event(
        self,
        event: Any,
        *,
        tts_publisher: AppGeneratorTTSPublisher | None = None,
        trace_manager: TraceQueueManager | None = None,
        queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None,
    ) -> Generator[StreamResponse]:
        """Dispatch events using elegant pattern matching."""
        handlers = self._get_event_handlers()
        event_type = type(event)

        # Direct handler lookup
        if handler := handlers.get(event_type):
            yield from handler(
                event,
                tts_publisher=tts_publisher,
                trace_manager=trace_manager,
                queue_message=queue_message,
            )
            return

        # Handle node failure events with isinstance check
        if isinstance(
            event,
            (
                QueueNodeFailedEvent,
                QueueNodeExceptionEvent,
            ),
        ):
            yield from self._handle_node_failed_events(
                event,
                tts_publisher=tts_publisher,
                trace_manager=trace_manager,
                queue_message=queue_message,
            )
            return

        # For unhandled events, we continue (original behavior)
        return

    def _process_stream_response(
        self,
        tts_publisher: AppGeneratorTTSPublisher | None = None,
        trace_manager: TraceQueueManager | None = None,
    ) -> Generator[StreamResponse]:
        """
        Process stream response using elegant Fluent Python patterns.
        Maintains exact same functionality as original 57-if-statement version.
        """
        for queue_message in self._base_task_pipeline.queue_manager.listen():
            event = queue_message.event

            match event:
                case QueueWorkflowStartedEvent():
                    self._resolve_graph_runtime_state()
                    yield from self._handle_workflow_started_event(event)

                case QueueErrorEvent():
                    yield from self._handle_error_event(event)
                    break

                case QueueWorkflowFailedEvent():
                    yield from self._handle_workflow_failed_event(event, trace_manager=trace_manager)
                    break
                case QueueWorkflowPausedEvent():
                    yield from self._handle_workflow_paused_event(event)
                    break

                case QueueWorkflowSucceededEvent():
                    yield from self._handle_workflow_succeeded_event(event, trace_manager=trace_manager)
                    break

                case QueueWorkflowPartialSuccessEvent():
                    yield from self._handle_workflow_partial_success_event(event, trace_manager=trace_manager)
                    break

                case QueueStopEvent():
                    yield from self._handle_stop_event(event, graph_runtime_state=None, trace_manager=trace_manager)
                    break

                # Handle all other events through elegant dispatch
                case _:
                    if responses := list(
                        self._dispatch_event(
                            event,
                            tts_publisher=tts_publisher,
                            trace_manager=trace_manager,
                            queue_message=queue_message,
                        )
                    ):
                        yield from responses

        if self._conversation_name_generate_thread:
            logger.debug("Conversation name generation running as daemon thread")

    def _save_message(self, *, session: Session, graph_runtime_state: RuntimeState | None = None):
        message = self._get_message(session=session)
        if message is None:
            return

        if message.status == MessageStatus.PAUSED:
            message.status = MessageStatus.NORMAL

        answer_text = self._task_state.answer

        message.answer = answer_text
        message.updated_at = naive_utc_now()
        message.provider_response_latency = time.perf_counter() - self._base_task_pipeline.start_at

        # Set usage first before dumping metadata
        if graph_runtime_state and graph_runtime_state.llm_usage:
            usage = graph_runtime_state.llm_usage
            message.message_tokens = usage.prompt_tokens
            message.message_unit_price = usage.prompt_unit_price
            message.message_price_unit = usage.prompt_price_unit
            message.answer_tokens = usage.completion_tokens
            message.answer_unit_price = usage.completion_unit_price
            message.answer_price_unit = usage.completion_price_unit
            message.total_price = usage.total_price
            message.currency = usage.currency
            self._task_state.metadata.usage = usage
        else:
            usage = LLMUsage.empty_usage()
            self._task_state.metadata.usage = usage

        # Add streaming metrics to usage if available
        if self._task_state.is_streaming_response and self._task_state.first_token_time:
            start_time = self._base_task_pipeline.start_at
            first_token_time = self._task_state.first_token_time
            last_token_time = self._task_state.last_token_time or first_token_time
            usage.time_to_first_token = round(first_token_time - start_time, 3)
            usage.time_to_generate = round(last_token_time - first_token_time, 3)

        metadata = self._task_state.metadata.model_dump()
        message.message_metadata = json.dumps(jsonable_encoder(metadata))
        message_files: list[MessageFile] = []
        for file in self._recorded_files:
            reference = file.get("reference") or file.get("related_id")
            message_files.append(
                MessageFile(
                    message_id=message.id,
                    type=file["type"],
                    transfer_method=file["transfer_method"],
                    url=file["remote_url"],
                    belongs_to=MessageFileBelongsTo.ASSISTANT,
                    upload_file_id=resolve_file_record_id(reference if isinstance(reference, str) else None),
                    created_by_role=CreatorUserRole.ACCOUNT
                    if message.invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}
                    else CreatorUserRole.END_USER,
                    created_by=message.from_account_id or message.from_end_user_id or "",
                )
            )
        session.add_all(message_files)

    def _emit_message_trace(self) -> None:
        trace_manager = self._application_generate_entity.trace_manager
        if trace_manager:
            trace_manager.add_trace_task(
                TraceTask(
                    TraceTaskName.MESSAGE_TRACE,
                    conversation_id=self._conversation_id,
                    message_id=self._message_id,
                    trace_session_id=self._application_generate_entity.extras.get("trace_session_id"),
                )
            )

    def _seed_graph_runtime_state_from_queue_manager(self) -> None:
        """Bootstrap the cached runtime state from the queue manager when present."""
        candidate = self._base_task_pipeline.queue_manager.graph_runtime_state
        if candidate is not None:
            self._graph_runtime_state = candidate

    def _message_end_to_stream_response(self) -> MessageEndStreamResponse:
        """
        Message end to stream response.
        :return:
        """
        extras = self._task_state.metadata.model_dump()

        if self._task_state.metadata.annotation_reply:
            del extras["annotation_reply"]

        return MessageEndStreamResponse(
            task_id=self._application_generate_entity.task_id,
            id=self._message_id,
            files=self._recorded_files,
            metadata=extras,
        )

    def _handle_output_moderation_chunk(self, text: str) -> bool:
        """
        Handle output moderation chunk.
        :param text: text
        :return: True if output moderation should direct output, otherwise False
        """
        if self._base_task_pipeline.output_moderation_handler:
            if self._base_task_pipeline.output_moderation_handler.should_direct_output():
                self._task_state.answer = self._base_task_pipeline.output_moderation_handler.get_final_output()
                self._base_task_pipeline.queue_manager.publish(
                    QueueTextChunkEvent(text=self._task_state.answer), PublishFrom.TASK_PIPELINE
                )

                self._base_task_pipeline.queue_manager.publish(
                    QueueStopEvent(stopped_by=QueueStopEvent.StopBy.OUTPUT_MODERATION), PublishFrom.TASK_PIPELINE
                )
                return True
            else:
                self._base_task_pipeline.output_moderation_handler.append_new_token(text)

        return False

    def _get_message(self, *, session: Session):
        stmt = select(Message).where(Message.id == self._message_id)
        message = session.scalar(stmt)
        if not message:
            raise ValueError(f"Message not found: {self._message_id}")
        return message

    def _save_output_for_event(self, event: QueueNodeSucceededEvent | QueueNodeExceptionEvent, node_execution_id: str):
        saver = self._draft_var_saver_factory(
            app_id=self._application_generate_entity.app_config.app_id,
            node_id=event.node_id,
            node_type=event.node_type,
            node_execution_id=node_execution_id,
            enclosing_node_id=event.in_loop_id or event.in_iteration_id,
        )
        saver.save(event.process_data, event.outputs)