dify/api/core/dify_builder/progress.py

"""Curated execution progress for long-running Builder handlers."""

import uuid
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from types import TracebackType
from typing import Literal

from core.dify_builder.contract import ExecutionActivity, ExecutionProgress, ProgressEventData
from core.dify_builder.models import NodeEvent, Session

__all__ = ["ProgressReporter"]

_InternalActivityState = Literal["pending", "active", "done", "failed", "stopped"]


@dataclass
class _ActivityRecord:
    id: str
    label: str
    state: _InternalActivityState = "pending"
    kind: Literal["stage", "node"] = "stage"
    parent_id: str | None = None


class ProgressReporter:
    """Publish ordered deltas for actions that have actually started.

    Handlers may declare their possible stages up front, but pending stages
    stay private. A stage enters the public stream only when activate is called.
    Workflow node events are folded into the same authoritative stream as child
    activities, so clients never have to merge two timelines. ``finish`` still
    returns a full snapshot for durable assistant output.
    """

    def __init__(
        self,
        *,
        emit: Callable[[ProgressEventData], None] | None,
        session_id: str,
        stage_id: str,
        at_version: int,
        activities: Iterable[ExecutionActivity],
        operation_id: str | None = None,
    ) -> None:
        self._emit = emit
        self._session_id = session_id
        self._stage_id = stage_id
        self._at_version = at_version
        self._operation_id = operation_id or str(uuid.uuid4())
        self._revision = 0
        self._status: Literal["running", "completed", "error", "stopped"] = "running"
        self._finished = False
        self._activities = [
            _ActivityRecord(
                id=activity.id,
                label=activity.label,
                kind=activity.kind,
                parent_id=activity.parent_id,
            )
            for activity in activities
        ]
        if len({activity.id for activity in self._activities}) != len(self._activities):
            raise ValueError("dify_builder: execution activity ids must be unique")

    @classmethod
    def for_session(
        cls,
        *,
        emit: Callable[[ProgressEventData], None] | None,
        session: Session,
        stage_id: str,
        steps: Iterable[tuple[str, str]],
        operation_id: str | None = None,
    ) -> "ProgressReporter":
        return cls(
            emit=emit,
            session_id=session.id,
            stage_id=stage_id,
            at_version=session.version + 1,
            activities=[ExecutionActivity(id=activity_id, label=label, state="active") for activity_id, label in steps],
            operation_id=operation_id,
        )

    @property
    def operation_id(self) -> str:
        return self._operation_id

    def __enter__(self) -> "ProgressReporter":
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        _exc: BaseException | None,
        _traceback: TracebackType | None,
    ) -> bool:
        if not self._finished:
            self.finish(status="error" if exc_type is not None else "completed")
        return False

    def activate(self, activity_id: str) -> None:
        """Complete prior active work and reveal the next stage."""
        activity = self._find(activity_id)
        self._status = "running"
        for current in self._activities:
            if current.state == "active" and current.id != activity_id:
                current.state = "done"
                self._publish(current)
        if activity.state != "active":
            activity.state = "active"
            self._publish(activity)

    def complete(self, activity_id: str) -> None:
        """Mark one revealed activity complete while the operation continues."""
        activity = self._find(activity_id)
        if activity.state == "done":
            return
        activity.state = "done"
        self._publish(activity)

    def add_steps(self, steps: Iterable[tuple[str, str]]) -> None:
        """Declare branch-specific stages without revealing them yet."""
        additions = [_ActivityRecord(id=activity_id, label=label) for activity_id, label in steps]
        known_ids = {activity.id for activity in self._activities}
        addition_ids = {activity.id for activity in additions}
        if len(addition_ids) != len(additions) or addition_ids & known_ids:
            raise ValueError("dify_builder: execution activity ids must be unique")
        self._activities.extend(additions)

    def fail_step(self, activity_id: str) -> None:
        """Mark one observable activity as failed before recovery continues."""
        activity = self._find(activity_id)
        if activity.state == "failed":
            return
        activity.state = "failed"
        self._publish(activity)

    def observe_node(self, parent_id: str, event: NodeEvent) -> None:
        """Fold a workflow node event into the canonical execution timeline."""
        activity_id = f"node:{event.node_id}"
        activity = next((item for item in self._activities if item.id == activity_id), None)
        if activity is None:
            activity = _ActivityRecord(
                id=activity_id,
                label=event.title or event.node_id,
                state="active",
                kind="node",
                parent_id=parent_id,
            )
            self._activities.append(activity)
        else:
            activity.label = event.title or event.node_id
            activity.parent_id = parent_id

        if event.status in {"success", "succeeded"}:
            activity.state = "done"
        elif event.status in {"error", "exception", "failed"}:
            activity.state = "failed"
        elif event.status == "running":
            activity.state = "active"
        else:
            activity.state = "stopped"
        self._publish(activity)

    def finish(
        self,
        *,
        status: Literal["completed", "error", "stopped"] = "completed",
    ) -> ExecutionProgress:
        """Publish terminal deltas and return the durable full snapshot."""
        if self._finished:
            return self._snapshot()
        active = [activity for activity in self._activities if activity.state == "active"]
        for activity in active:
            activity.state = "done" if status == "completed" else "failed" if status == "error" else "stopped"
        self._finished = True
        if not active:
            self._status = status
            self._publish(None)
            return self._snapshot()

        for activity in active[:-1]:
            self._publish(activity)
        self._status = status
        self._publish(active[-1])
        return self._snapshot()

    def _find(self, activity_id: str) -> _ActivityRecord:
        activity = next((item for item in self._activities if item.id == activity_id), None)
        if activity is None:
            raise ValueError(f"dify_builder: unknown execution activity {activity_id}")
        return activity

    def _snapshot(self) -> ExecutionProgress:
        return ExecutionProgress(
            status=self._status,
            activities=[self._as_activity(activity) for activity in self._activities if activity.state != "pending"],
        )

    @staticmethod
    def _as_activity(activity: _ActivityRecord) -> ExecutionActivity:
        assert activity.state != "pending", "pending execution activities are private"
        return ExecutionActivity(
            id=activity.id,
            label=activity.label,
            state=activity.state,
            kind=activity.kind,
            parent_id=activity.parent_id,
        )

    def _publish(self, activity: _ActivityRecord | None) -> None:
        self._revision += 1
        if self._emit is None:
            return
        self._emit(
            ProgressEventData(
                session_id=self._session_id,
                operation_id=self._operation_id,
                stage_id=self._stage_id,
                at_version=self._at_version,
                revision=self._revision,
                status=self._status,
                activity=self._as_activity(activity) if activity is not None else None,
            )
        )