dify/api/services/workflow_event_snapshot_service.py

from __future__ import annotations

import contextlib
import json
import logging
import queue
import threading
import time
from collections.abc import Generator, Mapping, Sequence
from dataclasses import dataclass, replace
from typing import Any

from sqlalchemy import desc, select
from sqlalchemy.orm import Session, sessionmaker

from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter
from core.app.apps.message_generator import MessageGenerator
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity
from core.app.entities.task_entities import (
    HumanInputRequiredResponse,
    MessageReplaceStreamResponse,
    NodeFinishStreamResponse,
    NodeStartStreamResponse,
    StreamEvent,
    WorkflowPauseStreamResponse,
    WorkflowStartStreamResponse,
)
from core.app.layers.pause_state_persist_layer import WorkflowResumptionContext
from core.workflow.human_input_forms import (
    load_form_dispositions_by_form_id,
)
from core.workflow.human_input_policy import (
    FormDisposition,
    HumanInputSurface,
    enrich_human_input_pause_reasons,
    resolve_human_input_pause_reason_inputs,
    resolve_variable_select_input_options,
)
from core.workflow.nodes.human_input.enums import HumanInputFormStatus
from core.workflow.nodes.human_input.pause_reason import (
    HumanInputRequired,
    PauseReason,
)
from graphon.entities import WorkflowStartReason
from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
from graphon.runtime import RuntimeState
from graphon.runtime.runtime_state_protocol import ReadOnlyVariablePool
from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter
from libs.broadcast_channel.exc import SubscriptionClosedError
from libs.datetime_utils import to_utc_timestamp
from models.human_input import HumanInputForm
from models.model import AppMode, Message
from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun
from repositories.api_workflow_node_execution_repository import WorkflowNodeExecutionSnapshot
from repositories.entities.workflow_pause import WorkflowPauseEntity
from repositories.factory import DifyAPIRepositoryFactory
from services.app_task_service import AppTaskService
from services.workflow_pause_service import WorkflowPauseSnapshot, load_workflow_pause_snapshot

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class MessageContext:
    conversation_id: str
    message_id: str
    created_at: int
    answer: str | None = None


@dataclass(frozen=True)
class _LegacyFormSnapshot:
    dispositions_by_form_id: Mapping[str, FormDisposition]
    expiration_times_by_form_id: Mapping[str, int]
    display_in_ui_by_form_id: Mapping[str, bool]


@dataclass
class BufferState:
    queue: queue.Queue[Mapping[str, Any]]
    stop_event: threading.Event
    done_event: threading.Event
    task_id_ready: threading.Event
    task_id_hint: str | None = None


def build_workflow_event_stream(
    *,
    app_mode: AppMode,
    workflow_run: WorkflowRun,
    tenant_id: str,
    app_id: str,
    session_maker: sessionmaker[Session],
    human_input_surface: HumanInputSurface | None = None,
    idle_timeout: float = 300,
    ping_interval: float = 10.0,
    close_on_pause: bool = True,
    include_node_details: bool = False,
) -> Generator[Mapping[str, Any] | str]:
    topic = MessageGenerator.get_response_topic(app_mode, workflow_run.id)
    workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
    node_execution_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(session_maker)

    pause_entity: WorkflowPauseEntity | None = None
    if workflow_run.status == WorkflowExecutionStatus.PAUSED:
        try:
            pause_entity = workflow_run_repo.get_workflow_pause(workflow_run.id)
        except Exception:
            logger.exception("Failed to load workflow pause for run %s", workflow_run.id)
            pause_entity = None

    resumption_context = _load_resumption_context(pause_entity)
    if (
        pause_entity is not None
        and resumption_context is not None
        and workflow_run.tenant_id == tenant_id
        and workflow_run.app_id == app_id
    ):
        AppTaskService.index_workflow_pause_task(
            workflow_run=workflow_run,
            pause_entity=pause_entity,
            resumption_context=resumption_context,
            workflow_run_repo=workflow_run_repo,
        )
    message_context: MessageContext | None = None
    if app_mode == AppMode.ADVANCED_CHAT:
        if workflow_run.status == WorkflowExecutionStatus.PAUSED:
            if resumption_context is None:
                raise AssertionError(
                    "WorkflowResumptionContext is required for advanced-chat snapshot replay, "
                    f"workflow_run_id={workflow_run.id}"
                )
            generate_entity = resumption_context.get_generate_entity()
            if not isinstance(generate_entity, AdvancedChatAppGenerateEntity):
                raise AssertionError(
                    "AdvancedChatAppGenerateEntity is required for advanced-chat snapshot replay, "
                    f"workflow_run_id={workflow_run.id}, generate_entity_type={type(generate_entity).__name__}"
                )
            if not generate_entity.conversation_id:
                raise AssertionError(
                    f"conversation_id is required for advanced-chat snapshot replay, workflow_run_id={workflow_run.id}"
                )
            message_context = _get_message_context_by_conversation(
                session_maker,
                conversation_id=generate_entity.conversation_id,
                workflow_run_id=workflow_run.id,
            )
        else:
            # Compatibility fallback for non-suspended snapshot requests. This app-scoped lookup is not optimal;
            # a dedicated index or stronger lookup key would be preferable.
            message_context = _get_message_context_by_app(
                session_maker,
                app_id=app_id,
                workflow_run_id=workflow_run.id,
            )

    node_snapshots = node_execution_repo.get_execution_snapshots_by_workflow_run(
        tenant_id=tenant_id,
        app_id=app_id,
        workflow_id=workflow_run.workflow_id,
        # NOTE(QuantumGhost): for events resumption, we only care about
        # the execution records from `WORKFLOW_RUN`.
        #
        # Ideally filtering with `workflow_run_id` is enough. However,
        # due to the index of `WorkflowNodeExecution` table, we have to
        # add a filter condition of `triggered_from` to
        # ensure that we can utilize the index.
        triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
        workflow_run_id=workflow_run.id,
    )
    # Console replay includes the same inline previews as run history. Other surfaces
    # retain the detail-free snapshot and avoid loading internal node data.
    node_executions = (
        {
            execution.node_execution_id or execution.id: execution
            for execution in node_execution_repo.get_executions_by_workflow_run(
                tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run.id
            )
        }
        if include_node_details
        else {}
    )

    def _generate() -> Generator[Mapping[str, Any] | str]:
        # send a PING event immediately to prevent the connection staying in pending state for a long time.
        #
        # This simplify the debugging process as the DevTools in Chrome does not
        # provide complete curl command for pending connections.
        yield StreamEvent.PING.value

        last_msg_time = time.time()
        last_ping_time = last_msg_time

        with topic.subscribe() as sub:
            buffer_state = _start_buffering(sub)
            try:
                task_id = _resolve_task_id(resumption_context, buffer_state, workflow_run.id)

                if workflow_run.status == WorkflowExecutionStatus.PAUSED and pause_entity is not None:
                    snapshot_events = _build_paused_snapshot_events(
                        workflow_run=workflow_run,
                        node_snapshots=node_snapshots,
                        task_id=task_id,
                        message_context=message_context,
                        node_executions=node_executions,
                        pause_entity=pause_entity,
                        resumption_context=resumption_context,
                        session_maker=session_maker,
                        human_input_surface=human_input_surface,
                    )
                else:
                    snapshot_events = _build_snapshot_events(
                        workflow_run=workflow_run,
                        node_snapshots=node_snapshots,
                        task_id=task_id,
                        message_context=message_context,
                        node_executions=node_executions,
                    )

                for event in snapshot_events:
                    last_msg_time = time.time()
                    last_ping_time = last_msg_time
                    yield event
                    if _is_terminal_event(event, close_on_pause=close_on_pause):
                        return

                while True:
                    if buffer_state.done_event.is_set() and buffer_state.queue.empty():
                        return

                    try:
                        event = buffer_state.queue.get(timeout=1)
                    except queue.Empty:
                        current_time = time.time()
                        if current_time - last_msg_time > idle_timeout:
                            logger.debug(
                                "Idle timeout of %s seconds reached, closing workflow event stream.",
                                idle_timeout,
                            )
                            return
                        if current_time - last_ping_time >= ping_interval:
                            yield StreamEvent.PING.value
                            last_ping_time = current_time
                        continue

                    last_msg_time = time.time()
                    last_ping_time = last_msg_time
                    yield event
                    if _is_terminal_event(event, close_on_pause=close_on_pause):
                        return
            finally:
                buffer_state.stop_event.set()

    return _generate()


def _get_message_context_by_conversation(
    session_maker: sessionmaker[Session],
    *,
    conversation_id: str,
    workflow_run_id: str,
) -> MessageContext | None:
    """Look up a paused or suspended Advanced Chat snapshot message by conversation and workflow run.

    Use this exact lookup after recovering ``conversation_id`` from persisted resumption context. Its predicates match
    ``message_workflow_run_id_idx``.
    """
    with session_maker() as session:
        stmt = (
            select(Message)
            .where(
                Message.conversation_id == conversation_id,
                Message.workflow_run_id == workflow_run_id,
            )
            .order_by(desc(Message.created_at))
            .limit(1)
        )
        message = session.scalar(stmt)
        if message is None:
            return None
        return _to_message_context(message)


def _get_message_context_by_app(
    session_maker: sessionmaker[Session],
    *,
    app_id: str,
    workflow_run_id: str,
) -> MessageContext | None:
    """Look up a non-suspended or running Advanced Chat reconnect snapshot by app and workflow run.

    This compatibility path applies only when no resumption context is expected. The app-scoped query is not optimal;
    a dedicated index or stronger lookup key would be preferable.
    """
    with session_maker() as session:
        stmt = (
            select(Message)
            .where(
                Message.app_id == app_id,
                Message.workflow_run_id == workflow_run_id,
            )
            .order_by(desc(Message.created_at))
            .limit(1)
        )
        message = session.scalar(stmt)
        if message is None:
            return None
        return _to_message_context(message)


def _to_message_context(message: Message) -> MessageContext:
    created_at = int(message.created_at.timestamp()) if message.created_at else 0
    return MessageContext(
        conversation_id=message.conversation_id,
        message_id=message.id,
        created_at=created_at,
        answer=message.answer,
    )


def _load_resumption_context(pause_entity: WorkflowPauseEntity | None) -> WorkflowResumptionContext | None:
    if pause_entity is None:
        return None
    try:
        raw_state = pause_entity.get_state().decode()
        return WorkflowResumptionContext.loads(raw_state)
    except Exception:
        logger.exception("Failed to load resumption context")
        return None


def _resolve_task_id(
    resumption_context: WorkflowResumptionContext | None,
    buffer_state: BufferState | None,
    workflow_run_id: str,
    wait_timeout: float = 0.2,
) -> str:
    if resumption_context is not None:
        generate_entity = resumption_context.get_generate_entity()
        if generate_entity.task_id:
            return generate_entity.task_id
    if buffer_state is None:
        return workflow_run_id
    if buffer_state.task_id_hint is None:
        buffer_state.task_id_ready.wait(timeout=wait_timeout)
    if buffer_state.task_id_hint:
        return buffer_state.task_id_hint
    return workflow_run_id


def _build_snapshot_events(
    *,
    workflow_run: WorkflowRun,
    node_snapshots: Sequence[WorkflowNodeExecutionSnapshot],
    task_id: str,
    message_context: MessageContext | None,
    node_executions: Mapping[str, WorkflowNodeExecutionModel],
) -> list[Mapping[str, Any]]:
    events: list[Mapping[str, Any]] = []

    workflow_started = _build_workflow_started_event(
        workflow_run=workflow_run,
        task_id=task_id,
    )
    _apply_message_context(workflow_started, message_context)
    events.append(workflow_started)

    if message_context is not None and message_context.answer is not None:
        message_replace = _build_message_replace_event(task_id=task_id, answer=message_context.answer)
        _apply_message_context(message_replace, message_context)
        events.append(message_replace)

    for snapshot in node_snapshots:
        execution = node_executions.get(snapshot.execution_id)
        node_started = _build_node_started_event(
            workflow_run_id=workflow_run.id,
            task_id=task_id,
            snapshot=snapshot,
            execution=execution,
        )
        _apply_message_context(node_started, message_context)
        events.append(node_started)

        if snapshot.status != WorkflowNodeExecutionStatus.RUNNING.value:
            node_finished = _build_node_finished_event(
                workflow_run_id=workflow_run.id,
                task_id=task_id,
                snapshot=snapshot,
                execution=execution,
            )
            _apply_message_context(node_finished, message_context)
            events.append(node_finished)

    return events


def _build_paused_snapshot_events(
    *,
    workflow_run: WorkflowRun,
    node_snapshots: Sequence[WorkflowNodeExecutionSnapshot],
    task_id: str,
    message_context: MessageContext | None,
    pause_entity: WorkflowPauseEntity,
    node_executions: Mapping[str, WorkflowNodeExecutionModel],
    resumption_context: WorkflowResumptionContext | None,
    session_maker: sessionmaker[Session],
    human_input_surface: HumanInputSurface | None = None,
) -> list[Mapping[str, Any]]:
    events = _build_snapshot_events(
        workflow_run=workflow_run,
        node_snapshots=node_snapshots,
        task_id=task_id,
        message_context=message_context,
        node_executions=node_executions,
    )
    pause_snapshot = load_workflow_pause_snapshot(
        pause_entity, session_factory=session_maker, context=resumption_context
    )
    open_legacy_ids = {
        reason.form_id
        for reason in _filter_open_pause_reasons(
            [
                reason
                for reason in pause_snapshot.reasons
                if isinstance(reason, HumanInputRequired) and reason.form_version == "1"
            ],
            session_maker=session_maker,
        )
        if isinstance(reason, HumanInputRequired)
    }
    pending_reasons = tuple(
        reason
        for reason in pause_snapshot.reasons
        if not isinstance(reason, HumanInputRequired)
        or (
            reason.form_id in open_legacy_ids
            if reason.form_version == "1"
            else pause_snapshot.v2_forms[reason.form_id].form.status == HumanInputFormStatus.WAITING
        )
    )
    if not pending_reasons:
        return events
    pause_snapshot = replace(pause_snapshot, reasons=pending_reasons)
    legacy_forms = _load_legacy_form_snapshot(
        pause_snapshot,
        session_maker=session_maker,
        human_input_surface=human_input_surface,
    )
    pause_events = _build_human_input_required_events(
        workflow_run_id=workflow_run.id,
        task_id=task_id,
        pause_snapshot=pause_snapshot,
        legacy_forms=legacy_forms,
    )
    pause_events.append(
        _build_pause_event(
            workflow_run=workflow_run,
            workflow_run_id=workflow_run.id,
            task_id=task_id,
            pause_snapshot=pause_snapshot,
            legacy_forms=legacy_forms,
        )
    )
    for event in pause_events:
        _apply_message_context(event, message_context)
        events.append(event)
    return events


def _build_workflow_started_event(
    *,
    workflow_run: WorkflowRun,
    task_id: str,
) -> dict[str, Any]:
    response = WorkflowStartStreamResponse(
        task_id=task_id,
        workflow_run_id=workflow_run.id,
        data=WorkflowStartStreamResponse.Data(
            id=workflow_run.id,
            workflow_id=workflow_run.workflow_id,
            inputs=workflow_run.inputs_dict or {},
            created_at=int(workflow_run.created_at.timestamp()),
            reason=WorkflowStartReason.INITIAL,
        ),
    )
    payload = response.model_dump(mode="json")
    payload["event"] = response.event.value
    return payload


def _build_message_replace_event(*, task_id: str, answer: str) -> dict[str, Any]:
    response = MessageReplaceStreamResponse(
        task_id=task_id,
        answer=answer,
        reason="",
    )
    payload = response.model_dump(mode="json")
    payload["event"] = response.event.value
    return payload


def _build_node_started_event(
    *,
    workflow_run_id: str,
    task_id: str,
    snapshot: WorkflowNodeExecutionSnapshot,
    execution: WorkflowNodeExecutionModel | None,
) -> dict[str, Any]:
    created_at = int(snapshot.created_at.timestamp()) if snapshot.created_at else 0
    response = NodeStartStreamResponse(
        task_id=task_id,
        workflow_run_id=workflow_run_id,
        data=NodeStartStreamResponse.Data(
            id=snapshot.execution_id,
            node_id=snapshot.node_id,
            node_type=snapshot.node_type,
            title=snapshot.title,
            index=snapshot.index,
            predecessor_node_id=execution.predecessor_node_id if execution else None,
            inputs=execution.inputs_dict if execution else None,
            inputs_truncated=execution.inputs_truncated if execution else False,
            created_at=created_at,
            extras=execution.extras if execution else {},
            iteration_id=snapshot.iteration_id,
            loop_id=snapshot.loop_id,
        ),
    )
    return response.model_dump(mode="json") if execution else response.to_ignore_detail_dict()


def _filter_open_pause_reasons(
    reasons: Sequence[PauseReason],
    *,
    session_maker: sessionmaker[Session] | None,
) -> list[PauseReason]:
    """Drop human-input reasons whose form is no longer waiting.

    Submitting a form only queues the resume, so the run stays paused until a worker picks it up.
    Replaying that form would ask again for an answer that was already given. Without a session the
    form status is unknown, so every reason is kept.
    """
    form_ids = [reason.form_id for reason in reasons if isinstance(reason, HumanInputRequired)]
    if not form_ids or session_maker is None:
        return list(reasons)

    stmt = select(HumanInputForm.id).where(
        HumanInputForm.id.in_(form_ids),
        HumanInputForm.status == HumanInputFormStatus.WAITING,
    )
    with session_maker() as session:
        waiting_form_ids = {str(form_id) for form_id in session.scalars(stmt)}
    return [
        reason for reason in reasons if not isinstance(reason, HumanInputRequired) or reason.form_id in waiting_form_ids
    ]


def _load_legacy_form_snapshot(
    pause_snapshot: WorkflowPauseSnapshot,
    *,
    session_maker: sessionmaker[Session],
    human_input_surface: HumanInputSurface | None,
) -> _LegacyFormSnapshot:
    human_input_form_ids = [
        reason.form_id
        for reason in pause_snapshot.reasons
        if isinstance(reason, HumanInputRequired) and reason.form_version == "1"
    ]

    expiration_times_by_form_id: dict[str, int] = {}
    display_in_ui_by_form_id: dict[str, bool] = {}
    dispositions_by_form_id: dict[str, FormDisposition] = {}
    if human_input_form_ids:
        stmt = select(HumanInputForm.id, HumanInputForm.expiration_time, HumanInputForm.form_definition).where(
            HumanInputForm.id.in_(human_input_form_ids)
        )
        with session_maker() as session:
            for form_id, expiration_time, form_definition in session.execute(stmt):
                expiration_times_by_form_id[str(form_id)] = to_utc_timestamp(expiration_time)
                try:
                    definition_payload = json.loads(form_definition) if form_definition else {}
                except (TypeError, json.JSONDecodeError):
                    definition_payload = {}
                display_in_ui_by_form_id[str(form_id)] = bool(definition_payload.get("display_in_ui"))
            dispositions_by_form_id = load_form_dispositions_by_form_id(
                human_input_form_ids,
                session=session,
                surface=human_input_surface,
            )

    return _LegacyFormSnapshot(dispositions_by_form_id, expiration_times_by_form_id, display_in_ui_by_form_id)


def _build_human_input_required_events(
    *,
    workflow_run_id: str,
    task_id: str,
    pause_snapshot: WorkflowPauseSnapshot,
    legacy_forms: _LegacyFormSnapshot,
) -> list[dict[str, Any]]:
    reasons = pause_snapshot.reasons
    v2_forms = pause_snapshot.v2_forms
    variable_pool = pause_snapshot.runtime_state.variable_pool
    events: list[dict[str, Any]] = []
    for reason in reasons:
        if not isinstance(reason, HumanInputRequired):
            continue

        if reason.form_version == "2":
            prepared = v2_forms.get(reason.form_id)
            if prepared is not None:
                response = HumanInputRequiredResponse(
                    task_id=task_id,
                    workflow_run_id=workflow_run_id,
                    data=WorkflowResponseConverter.human_input_v2_required_data(prepared, node_id=reason.node_id),
                )
                events.append(response.model_dump(mode="json"))
            continue

        form_id = reason.form_id

        expiration_time = legacy_forms.expiration_times_by_form_id.get(form_id)
        if expiration_time is None:
            continue

        resolved_inputs = resolve_variable_select_input_options(
            reason.inputs,
            variable_pool=variable_pool,
        )
        disposition = legacy_forms.dispositions_by_form_id.get(form_id)

        response = HumanInputRequiredResponse(
            task_id=task_id,
            workflow_run_id=workflow_run_id,
            data=HumanInputRequiredResponse.Data(
                form_version=reason.form_version,
                form_id=form_id,
                node_id=reason.node_id,
                node_title=reason.node_title,
                form_content=reason.form_content,
                inputs=resolved_inputs,
                actions=reason.actions,
                display_in_ui=legacy_forms.display_in_ui_by_form_id.get(form_id, False),
                form_token=disposition.form_token if disposition else None,
                approval_channels=list(disposition.approval_channels) if disposition else [],
                resolved_default_values=reason.resolved_default_values,
                expiration_time=expiration_time,
            ),
        )
        payload = response.model_dump(mode="json")
        payload["event"] = response.event.value
        events.append(payload)

    return events


def _load_variable_pool_from_resumption_context(
    resumption_context: WorkflowResumptionContext | None,
) -> ReadOnlyVariablePool | None:
    if resumption_context is None:
        return None
    state = RuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state)

    return state.variable_pool


def _build_node_finished_event(
    *,
    workflow_run_id: str,
    task_id: str,
    snapshot: WorkflowNodeExecutionSnapshot,
    execution: WorkflowNodeExecutionModel | None,
) -> dict[str, Any]:
    created_at = int(snapshot.created_at.timestamp()) if snapshot.created_at else 0
    finished_at = int(snapshot.finished_at.timestamp()) if snapshot.finished_at else created_at
    response = NodeFinishStreamResponse(
        task_id=task_id,
        workflow_run_id=workflow_run_id,
        data=NodeFinishStreamResponse.Data(
            id=snapshot.execution_id,
            node_id=snapshot.node_id,
            node_type=snapshot.node_type,
            title=snapshot.title,
            index=snapshot.index,
            predecessor_node_id=execution.predecessor_node_id if execution else None,
            inputs=execution.inputs_dict if execution else None,
            inputs_truncated=execution.inputs_truncated if execution else False,
            process_data=execution.process_data_dict if execution else None,
            process_data_truncated=execution.process_data_truncated if execution else False,
            outputs=execution.outputs_dict if execution else None,
            outputs_truncated=execution.outputs_truncated if execution else False,
            status=WorkflowNodeExecutionStatus(snapshot.status),
            error=execution.error if execution else None,
            elapsed_time=snapshot.elapsed_time,
            execution_metadata=(
                {
                    WorkflowNodeExecutionMetadataKey(key): value
                    for key, value in execution.execution_metadata_dict.items()
                }
                if execution
                else None
            ),
            created_at=created_at,
            finished_at=finished_at,
            files=[],
            iteration_id=snapshot.iteration_id,
            loop_id=snapshot.loop_id,
        ),
    )
    return response.model_dump(mode="json") if execution else response.to_ignore_detail_dict()


def _build_pause_event(
    *,
    workflow_run: WorkflowRun,
    workflow_run_id: str,
    task_id: str,
    pause_snapshot: WorkflowPauseSnapshot,
    legacy_forms: _LegacyFormSnapshot,
) -> dict[str, Any]:
    state = pause_snapshot.runtime_state
    outputs = dict(WorkflowRuntimeTypeConverter().to_json_encodable(state.outputs or {}))
    variable_pool = state.variable_pool

    resolved_pause_reasons = resolve_human_input_pause_reason_inputs(
        pause_snapshot.reasons,
        variable_pool=variable_pool,
    )
    paused_nodes = list(
        dict.fromkeys(reason.node_id for reason in resolved_pause_reasons if isinstance(reason, HumanInputRequired))
    )
    reasons = [reason.model_dump(mode="json") for reason in resolved_pause_reasons]
    human_input_form_ids = [
        reason.form_id
        for reason in resolved_pause_reasons
        if isinstance(reason, HumanInputRequired) and reason.form_version == "1"
    ]
    if human_input_form_ids:
        # Reconnect paths must preserve the same pause-reason contract as live streams;
        # otherwise clients see schema drift after resume.
        reasons = enrich_human_input_pause_reasons(
            reasons,
            dispositions_by_form_id=legacy_forms.dispositions_by_form_id,
            expiration_times_by_form_id=legacy_forms.expiration_times_by_form_id,
        )

    v2_forms = pause_snapshot.v2_forms
    for reason in resolved_pause_reasons:
        if not isinstance(reason, HumanInputRequired) or reason.form_version != "2":
            continue
        prepared = v2_forms.get(reason.form_id)
        if prepared is None:
            continue
        presentation = WorkflowResponseConverter.human_input_v2_required_data(prepared, node_id=reason.node_id)
        for payload in reasons:
            if payload.get("form_id") == reason.form_id:
                payload.update(presentation.model_dump(mode="json"))

    response = WorkflowPauseStreamResponse(
        task_id=task_id,
        workflow_run_id=workflow_run_id,
        data=WorkflowPauseStreamResponse.Data(
            workflow_run_id=workflow_run_id,
            paused_nodes=paused_nodes,
            outputs=outputs,
            reasons=reasons,
            status=workflow_run.status,
            created_at=int(workflow_run.created_at.timestamp()),
            elapsed_time=float(workflow_run.elapsed_time or 0.0),
            total_tokens=int(workflow_run.total_tokens or 0),
            total_steps=int(workflow_run.total_steps or 0),
        ),
    )
    payload = response.model_dump(mode="json")
    payload["event"] = response.event.value
    return payload


def _apply_message_context(payload: dict[str, Any], message_context: MessageContext | None) -> None:
    if message_context is None:
        return
    payload["conversation_id"] = message_context.conversation_id
    payload["message_id"] = message_context.message_id
    payload["created_at"] = message_context.created_at


def _start_buffering(subscription) -> BufferState:
    buffer_state = BufferState(
        queue=queue.Queue(maxsize=2048),
        stop_event=threading.Event(),
        done_event=threading.Event(),
        task_id_ready=threading.Event(),
    )

    def _worker() -> None:
        dropped_count = 0
        try:
            while not buffer_state.stop_event.is_set():
                msg = subscription.receive(timeout=1)
                if msg is None:
                    continue
                event = _parse_event_message(msg)
                if event is None:
                    continue
                task_id = event.get("task_id")
                if task_id and buffer_state.task_id_hint is None:
                    buffer_state.task_id_hint = str(task_id)
                    buffer_state.task_id_ready.set()
                try:
                    buffer_state.queue.put_nowait(event)
                except queue.Full:
                    dropped_count += 1
                    with contextlib.suppress(queue.Empty):
                        buffer_state.queue.get_nowait()
                    try:
                        buffer_state.queue.put_nowait(event)
                    except queue.Full:
                        continue
                    logger.warning("Dropped buffered workflow event, total_dropped=%s", dropped_count)
        except SubscriptionClosedError:
            pass
        except Exception:
            logger.exception("Failed while buffering workflow events")
        finally:
            buffer_state.done_event.set()

    thread = threading.Thread(target=_worker, name=f"workflow-event-buffer-{id(subscription)}", daemon=True)
    thread.start()
    return buffer_state


def _parse_event_message(message: bytes) -> Mapping[str, Any] | None:
    try:
        event = json.loads(message)
    except json.JSONDecodeError:
        logger.warning("Failed to decode workflow event payload")
        return None
    if not isinstance(event, dict):
        return None
    return event


def _is_terminal_event(
    event: Mapping[str, Any] | str,
    close_on_pause: bool = True,
    *,
    include_paused: bool | None = None,
) -> bool:
    if include_paused is not None:
        close_on_pause = include_paused
    if not isinstance(event, Mapping):
        return False
    event_type = event.get("event")
    if event_type == StreamEvent.WORKFLOW_FINISHED.value:
        return True
    if close_on_pause:
        return event_type == StreamEvent.WORKFLOW_PAUSED.value
    return False