ai-video-fullstack/backend/services/action_runtime.py

"""Deterministic Action execution shared by Prompt startup and Workflow."""

from __future__ import annotations

import asyncio
from collections.abc import Callable
from dataclasses import dataclass
from enum import StrEnum
from time import monotonic
from typing import Any
from uuid import uuid4

from models import RuntimeTool
from services.tool_executor import ToolExecutionError, ToolExecutor


class ActionStatus(StrEnum):
    SUCCESS = "success"
    FAILURE = "failure"
    CANCELLED = "cancelled"


@dataclass(frozen=True)
class ActionError:
    code: str
    message: str
    retryable: bool = False


@dataclass(frozen=True)
class ActionOutcome:
    """One completed Action invocation.

    Raw tool results stay in memory. Persisted or client-visible events should
    use ``trace_payload`` so private business data is not copied accidentally.
    """

    invocation_id: str
    status: ActionStatus
    duration_ms: int
    result: dict[str, Any] | None = None
    updated_variables: tuple[str, ...] = ()
    error: ActionError | None = None

    @property
    def should_route(self) -> bool:
        return self.status != ActionStatus.CANCELLED

    def trace_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "invocationId": self.invocation_id,
            "status": self.status.value,
            "durationMs": self.duration_ms,
            "updatedVariables": list(self.updated_variables),
        }
        if self.result is not None:
            payload["resultKeys"] = sorted(str(key) for key in self.result)
        if self.error is not None:
            payload["error"] = {
                "code": self.error.code,
                "message": self.error.message,
                "retryable": self.error.retryable,
            }
        return payload


class ActionInvocationCancelled(asyncio.CancelledError):
    """Carry a structured outcome while preserving task cancellation."""

    def __init__(self, outcome: ActionOutcome) -> None:
        super().__init__(outcome.error.message if outcome.error else "Action cancelled")
        self.outcome = outcome


class ActionRunner:
    """Normalize ToolExecutor's heterogeneous responses into ActionOutcome."""

    _CANCELLED_MARKERS = (
        "会话已结束",
        "会话已取消",
        "管线已停止",
        "通道已关闭",
        "连接已断开",
    )

    def __init__(
        self,
        executor: ToolExecutor,
        *,
        is_session_ending: Callable[[], bool] | None = None,
    ) -> None:
        self._executor = executor
        self._is_session_ending = is_session_ending or (lambda: False)

    @staticmethod
    def new_invocation_id() -> str:
        return f"act_{uuid4().hex[:20]}"

    async def execute(
        self,
        tool: RuntimeTool | None,
        arguments: dict[str, Any] | None = None,
        *,
        result_assignments: dict[str, str] | None = None,
        invocation_id: str | None = None,
    ) -> ActionOutcome:
        invocation_id = invocation_id or self.new_invocation_id()
        started_at = monotonic()
        result: dict[str, Any] | None = None
        if tool is None:
            return self._failure(
                invocation_id,
                started_at,
                code="tool_not_found",
                message="Action 引用的工具不存在",
            )

        try:
            rendered_arguments = self._executor.store.render_data(arguments or {})
            result = await self._executor.execute(
                tool,
                rendered_arguments,
                result_assignments=result_assignments,
            )
            if result.get("status") != "ok":
                returned_status = str(result.get("status") or "error")
                message = str(result.get("message") or "工具返回执行失败状态")
                if self._is_cancelled(message):
                    return ActionOutcome(
                        invocation_id=invocation_id,
                        status=ActionStatus.CANCELLED,
                        duration_ms=self._elapsed_ms(started_at),
                        result=result,
                        error=ActionError(
                            code="session_ended",
                            message=message[:2048],
                        ),
                    )
                return self._failure(
                    invocation_id,
                    started_at,
                    code=str(result.get("code") or f"tool_{returned_status}"),
                    message=message,
                    retryable=bool(result.get("retryable", False)),
                    result=result,
                )
            return ActionOutcome(
                invocation_id=invocation_id,
                status=ActionStatus.SUCCESS,
                duration_ms=self._elapsed_ms(started_at),
                result=result,
                updated_variables=tuple(
                    str(name) for name in result.get("updated_variables") or []
                ),
            )
        except asyncio.CancelledError as exc:
            outcome = ActionOutcome(
                invocation_id=invocation_id,
                status=ActionStatus.CANCELLED,
                duration_ms=self._elapsed_ms(started_at),
                result=result,
                error=ActionError(
                    code="action_cancelled",
                    message="Action 随当前任务取消",
                ),
            )
            raise ActionInvocationCancelled(outcome) from exc
        except (ToolExecutionError, ValueError) as exc:
            cancelled = self._is_cancelled(str(exc))
            if cancelled:
                return ActionOutcome(
                    invocation_id=invocation_id,
                    status=ActionStatus.CANCELLED,
                    duration_ms=self._elapsed_ms(started_at),
                    result=result,
                    error=ActionError(
                        code="session_ended",
                        message=str(exc)[:2048],
                    ),
                )
            return self._failure(
                invocation_id,
                started_at,
                code=(
                    "invalid_action_configuration"
                    if isinstance(exc, ValueError)
                    else "tool_execution_error"
                ),
                message=str(exc),
                result=result,
            )

    def _is_cancelled(self, message: str) -> bool:
        return self._is_session_ending() or any(
            marker in message for marker in self._CANCELLED_MARKERS
        )

    def _failure(
        self,
        invocation_id: str,
        started_at: float,
        *,
        code: str,
        message: str,
        retryable: bool = False,
        result: dict[str, Any] | None = None,
    ) -> ActionOutcome:
        return ActionOutcome(
            invocation_id=invocation_id,
            status=ActionStatus.FAILURE,
            duration_ms=self._elapsed_ms(started_at),
            result=result,
            error=ActionError(
                code=code,
                message=message[:2048],
                retryable=retryable,
            ),
        )

    @staticmethod
    def _elapsed_ms(started_at: float) -> int:
        return max(0, round((monotonic() - started_at) * 1000))