import asyncio
import copy
import functools
import inspect
import json
import sys
import threading
import time
import uuid
from concurrent.futures import Future
from functools import wraps
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Coroutine,
    Generic,
    List,
    Optional,
    ParamSpec,
    TypedDict,
    TypeVar,
    Union,
    cast,
)

from dbos._outcome import NoResult, Outcome, Pending
from dbos._utils import GlobalParams, retriable_postgres_exception

from ._app_db import ApplicationDatabase, TransactionResultInternal
from ._context import (
    DBOSAssumeRole,
    DBOSContext,
    DBOSContextEnsure,
    DBOSContextSetAuth,
    EnterDBOSStepCtx,
    EnterDBOSTransaction,
    EnterDBOSWorkflow,
    OperationType,
    SetWorkflowID,
    TracedAttributes,
    assert_current_dbos_context,
    get_local_dbos_context,
)
from ._error import (
    DBOSAwaitedWorkflowCancelledError,
    DBOSException,
    DBOSMaxStepRetriesExceeded,
    DBOSNonExistentWorkflowError,
    DBOSQueueDeduplicatedError,
    DBOSRecoveryError,
    DBOSUnexpectedStepError,
    DBOSWorkflowCancelledError,
    DBOSWorkflowConflictIDError,
    DBOSWorkflowFunctionNotFoundError,
)
from ._logger import dbos_logger
from ._registrations import (
    DEFAULT_MAX_RECOVERY_ATTEMPTS,
    DBOSFuncType,
    ValidateArgsCallable,
    get_config_name,
    get_dbos_class_name,
    get_dbos_func_name,
    get_func_info,
    get_or_create_func_info,
    set_dbos_func_name,
    set_func_info,
    set_temp_workflow_type,
)
from ._roles import check_required_roles
from ._serialization import (
    DBOSPortableJSON,
    WorkflowInputs,
    WorkflowSerializationFormat,
    deserialize_args,
    deserialize_exception,
    deserialize_value,
    serialize_args,
    serialize_exception,
    serialize_value,
    serialize_value_as,
)
from ._sys_db import (
    EnqueueOptionsInternal,
    OperationResultInternal,
    WorkflowStatus,
    WorkflowStatusInternal,
    WorkflowStatusString,
)

if TYPE_CHECKING:
    from ._dbos import (
        DBOS,
        WorkflowHandle,
        WorkflowHandleAsync,
        DBOSRegistry,
        IsolationLevel,
    )

from sqlalchemy.exc import DBAPIError, InvalidRequestError

P = ParamSpec("P")  # A generic type for workflow parameters
R = TypeVar("R", covariant=True)  # A generic type for workflow return values
F = TypeVar("F", bound=Callable[..., Any])

TEMP_SEND_WF_NAME = "<temp>.temp_send_workflow"
DEBOUNCER_WORKFLOW_NAME = "_dbos_debouncer_workflow"
DEFAULT_POLLING_INTERVAL = 1.0


class WorkflowHandleFuture(Generic[R]):

    def __init__(self, workflow_id: str, future: Future[R], dbos: "DBOS"):
        self.workflow_id = workflow_id
        self.future = future
        self.dbos = dbos

    def get_workflow_id(self) -> str:
        return self.workflow_id

    def get_result(
        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
    ) -> R:
        try:
            r = self.future.result()
        except Exception as e:
            serialized_e, serialization = serialize_exception(
                e, None, self.dbos._serializer
            )
            self.dbos._sys_db.record_get_result(
                self.workflow_id, None, serialized_e, serialization
            )
            raise
        serialized_r, serialization = serialize_value(r, None, self.dbos._serializer)
        self.dbos._sys_db.record_get_result(
            self.workflow_id, serialized_r, None, serialization
        )
        return r

    def get_status(self) -> WorkflowStatus:
        stat = self.dbos.get_workflow_status(self.workflow_id)
        if stat is None:
            raise DBOSNonExistentWorkflowError("target", self.workflow_id)
        return stat


class WorkflowHandlePolling(Generic[R]):

    def __init__(self, workflow_id: str, dbos: "DBOS"):
        self.workflow_id = workflow_id
        self.dbos = dbos

    def get_workflow_id(self) -> str:
        return self.workflow_id

    def get_result(
        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
    ) -> R:
        try:
            r: R = self.dbos._sys_db.await_workflow_result(
                self.workflow_id, polling_interval_sec
            )
        except Exception as e:
            serialized_e, serialization = serialize_exception(
                e, None, self.dbos._serializer
            )
            self.dbos._sys_db.record_get_result(
                self.workflow_id, None, serialized_e, serialization
            )
            raise
        serialized_r, serialization = serialize_value(r, None, self.dbos._serializer)
        self.dbos._sys_db.record_get_result(
            self.workflow_id, serialized_r, None, serialization
        )
        return r

    def get_status(self) -> WorkflowStatus:
        stat = self.dbos.get_workflow_status(self.workflow_id)
        if stat is None:
            raise DBOSNonExistentWorkflowError("target", self.workflow_id)
        return stat


class WorkflowHandleAsyncTask(Generic[R]):

    def __init__(self, workflow_id: str, task: asyncio.Future[R], dbos: "DBOS"):
        self.workflow_id = workflow_id
        self.task = task
        self.dbos = dbos

    def get_workflow_id(self) -> str:
        return self.workflow_id

    async def get_result(
        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
    ) -> R:
        try:
            r = await self.task
        except Exception as e:
            serialized_e, serialization = serialize_exception(
                e, None, self.dbos._serializer
            )
            await asyncio.to_thread(
                self.dbos._sys_db.record_get_result,
                self.workflow_id,
                None,
                serialized_e,
                serialization,
            )
            raise
        serialized_r, serialization = serialize_value(r, None, self.dbos._serializer)
        await asyncio.to_thread(
            self.dbos._sys_db.record_get_result,
            self.workflow_id,
            serialized_r,
            None,
            serialization,
        )
        return r

    async def get_status(self) -> WorkflowStatus:
        stat = await asyncio.to_thread(self.dbos.get_workflow_status, self.workflow_id)
        if stat is None:
            raise DBOSNonExistentWorkflowError("target", self.workflow_id)
        return stat


class WorkflowHandleAsyncPolling(Generic[R]):

    def __init__(self, workflow_id: str, dbos: "DBOS"):
        self.workflow_id = workflow_id
        self.dbos = dbos

    def get_workflow_id(self) -> str:
        return self.workflow_id

    async def get_result(
        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
    ) -> R:
        try:
            r: R = await self.dbos._sys_db.await_workflow_result_async(
                self.workflow_id,
                polling_interval_sec,
            )
        except Exception as e:
            serialized_e, serialization = serialize_exception(
                e, None, self.dbos._serializer
            )
            await asyncio.to_thread(
                self.dbos._sys_db.record_get_result,
                self.workflow_id,
                None,
                serialized_e,
                serialization,
            )
            raise
        serialized_r, serialization = serialize_value(r, None, self.dbos._serializer)
        await asyncio.to_thread(
            self.dbos._sys_db.record_get_result,
            self.workflow_id,
            serialized_r,
            None,
            serialization,
        )
        return r

    async def get_status(self) -> WorkflowStatus:
        stat = await asyncio.to_thread(self.dbos.get_workflow_status, self.workflow_id)
        if stat is None:
            raise DBOSNonExistentWorkflowError("target", self.workflow_id)
        return stat


class StepOptions(TypedDict, total=False):
    """
    Configuration options for steps.

    Attributes:
        name:
            Optional name for the step.
            If not provided, the function's name will be used.

        retries_allowed:
            Whether the step should be retried on failure.

        interval_seconds:
            Initial delay (in seconds) between retry attempts.

        max_attempts:
            Maximum number of attempts before the step is
            considered failed.

        backoff_rate:
            Multiplier applied to `interval_seconds` after
            each failed attempt (e.g. 2.0 = exponential backoff).
    """

    name: Optional[str]
    retries_allowed: bool
    interval_seconds: float
    max_attempts: int
    backoff_rate: float


DEFAULT_STEP_OPTIONS: StepOptions = {
    "name": None,
    "retries_allowed": False,
    "interval_seconds": 1.0,
    "max_attempts": 3,
    "backoff_rate": 2.0,
}


def normalize_step_options(opts: Optional[StepOptions]) -> StepOptions:
    return {**DEFAULT_STEP_OPTIONS, **(opts or {})}


def _init_workflow(
    dbos: "DBOS",
    ctx: DBOSContext,
    *,
    inputs: WorkflowInputs,
    wf_name: str,
    class_name: Optional[str],
    config_name: Optional[str],
    queue: Optional[str],
    workflow_timeout_ms: Optional[int],
    workflow_deadline_epoch_ms: Optional[int],
    max_recovery_attempts: Optional[int],
    enqueue_options: Optional[EnqueueOptionsInternal],
    is_recovery_request: Optional[bool],
    is_dequeued_request: Optional[bool],
    serialization_type: Optional[WorkflowSerializationFormat],
) -> tuple[WorkflowStatusInternal, bool]:
    wfid = (
        ctx.workflow_id
        if len(ctx.workflow_id) > 0
        else ctx.id_assigned_for_next_workflow
    )

    # If we have a class name, the first arg is the instance and do not serialize
    if class_name is not None and class_name != "":
        inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}

    sertype: WorkflowSerializationFormat | None = serialization_type
    if sertype is None or sertype == WorkflowSerializationFormat.DEFAULT:
        sertype = ctx.serialization_type
    serargs, serialization = serialize_args(
        inputs["args"], inputs["kwargs"], sertype, dbos._serializer
    )

    # Initialize a workflow status object from the context
    status: WorkflowStatusInternal = {
        "workflow_uuid": wfid,
        "status": (
            WorkflowStatusString.PENDING.value
            if queue is None
            else WorkflowStatusString.ENQUEUED.value
        ),
        "name": wf_name,
        "class_name": class_name,
        "config_name": config_name,
        "output": None,
        "error": None,
        "app_id": ctx.app_id,
        "app_version": (
            enqueue_options["app_version"]
            if enqueue_options is not None
            and enqueue_options["app_version"] is not None
            else GlobalParams.app_version
        ),
        "executor_id": ctx.executor_id,
        "recovery_attempts": None,
        "authenticated_user": ctx.authenticated_user,
        "authenticated_roles": (
            json.dumps(ctx.authenticated_roles) if ctx.authenticated_roles else None
        ),
        "assumed_role": ctx.assumed_role,
        "queue_name": queue,
        "created_at": None,
        "updated_at": None,
        "workflow_timeout_ms": workflow_timeout_ms,
        "workflow_deadline_epoch_ms": workflow_deadline_epoch_ms,
        "deduplication_id": (
            enqueue_options["deduplication_id"] if enqueue_options is not None else None
        ),
        "priority": (
            (
                enqueue_options["priority"]
                if enqueue_options["priority"] is not None
                else 0
            )
            if enqueue_options is not None
            else 0
        ),
        "inputs": serargs,
        "queue_partition_key": (
            enqueue_options["queue_partition_key"]
            if enqueue_options is not None
            else None
        ),
        "forked_from": None,
        "parent_workflow_id": (
            ctx.parent_workflow_id if len(ctx.parent_workflow_id) > 0 else None
        ),
        "started_at_epoch_ms": None,
        "owner_xid": None,
        "serialization": serialization,
    }

    # Synchronously record the status and inputs for workflows
    try:
        wf_status, workflow_deadline_epoch_ms, should_execute = (
            dbos._sys_db.init_workflow(
                status,
                max_recovery_attempts=max_recovery_attempts,
                is_dequeued_request=is_dequeued_request,
                is_recovery_request=is_recovery_request,
                owner_xid=str(uuid.uuid4()),
            )
        )
    except DBOSQueueDeduplicatedError as e:
        sererr, serialization = serialize_exception(
            e,
            status["serialization"],
            dbos._serializer,
        )
        if ctx.has_parent():
            result: OperationResultInternal = {
                "workflow_uuid": ctx.parent_workflow_id,
                "function_id": ctx.parent_workflow_fid,
                "function_name": wf_name,
                "output": None,
                "error": sererr,
                "serialization": serialization,
                "started_at_epoch_ms": int(time.time() * 1000),
            }
            dbos._sys_db.record_operation_result(result)
        raise

    if should_execute and workflow_deadline_epoch_ms is not None:
        evt = threading.Event()
        dbos.background_thread_stop_events.append(evt)

        def timeout_func() -> None:
            try:
                assert workflow_deadline_epoch_ms is not None
                time_to_wait_sec = (
                    workflow_deadline_epoch_ms - (time.time() * 1000)
                ) / 1000
                if time_to_wait_sec > 0:
                    was_stopped = evt.wait(time_to_wait_sec)
                    if was_stopped:
                        return
                dbos._sys_db.cancel_workflow(wfid)
            except Exception as e:
                dbos.logger.warning(
                    f"Exception in timeout thread for workflow {wfid}: {e}"
                )

        timeout_thread = threading.Thread(target=timeout_func, daemon=True)
        timeout_thread.start()
        dbos._background_threads.append(timeout_thread)

    ctx.workflow_deadline_epoch_ms = workflow_deadline_epoch_ms
    status["workflow_deadline_epoch_ms"] = workflow_deadline_epoch_ms
    status["status"] = wf_status
    return status, should_execute


def _get_wf_invoke_func(
    dbos: "DBOS",
    status: WorkflowStatusInternal,
) -> Callable[[Callable[[], R]], R]:
    def persist(func: Callable[[], R]) -> R:
        if (
            status["status"] == WorkflowStatusString.ERROR.value
            or status["status"] == WorkflowStatusString.SUCCESS.value
        ):
            dbos.logger.debug(
                f"Workflow {status['workflow_uuid']} is already completed with status {status['status']}"
            )
            # Directly return the result if the workflow is already completed
            recorded_result: R = dbos._sys_db.await_workflow_result(
                status["workflow_uuid"], polling_interval=DEFAULT_POLLING_INTERVAL
            )
            return recorded_result
        try:
            if inspect.iscoroutinefunction(func):
                output = dbos._background_event_loop.submit_coroutine(
                    cast(Coroutine[Any, Any, R], func())
                )
            else:
                output = func()

            serval, _serialization = serialize_value_as(
                output, status["serialization"], dbos._serializer
            )
            dbos._sys_db.update_workflow_outcome(
                status["workflow_uuid"],
                WorkflowStatusString.SUCCESS.value,
                output=serval,
            )
            return output
        except DBOSWorkflowConflictIDError:
            # Await the workflow result
            r: R = dbos._sys_db.await_workflow_result(
                status["workflow_uuid"], polling_interval=DEFAULT_POLLING_INTERVAL
            )
            return r
        except DBOSWorkflowCancelledError as error:
            raise DBOSAwaitedWorkflowCancelledError(status["workflow_uuid"])
        except Exception as error:
            dbos._sys_db.update_workflow_outcome(
                status["workflow_uuid"],
                WorkflowStatusString.ERROR.value,
                error=serialize_exception(
                    error,
                    status["serialization"],
                    dbos._serializer,
                )[0],
            )
            raise

    return persist


class ActiveWorkflowById:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._m: dict[str, bool] = {}

    def acquire(self, key: str) -> bool:
        """
        Returns is_owner
        """
        with self._lock:
            entry = self._m.get(key)
            if entry is None:
                self._m[key] = True
                return True
            return False

    def release(
        self,
        key: str,
    ) -> None:
        """
        Removes the key when work done
        """
        with self._lock:
            del self._m[key]

    def activeList(self) -> List[str]:
        return list(self._m.keys())


def _execute_workflow_wthread(
    dbos: "DBOS",
    status: WorkflowStatusInternal,
    func: "Callable[P, R]",
    ctx: DBOSContext,
    args: tuple[Any],
    kwargs: dict[str, Any],
) -> R:
    attributes: TracedAttributes = {
        "name": get_dbos_func_name(func),
        "operationType": OperationType.WORKFLOW.value,
    }
    fi = get_func_info(func)
    with EnterDBOSWorkflow(attributes, ctx):
        rr: Optional[str] = check_required_roles(func, fi)
        with DBOSAssumeRole(rr):
            owned = dbos._active_workflows_set.acquire(status["workflow_uuid"])
            try:
                if owned:
                    return _get_wf_invoke_func(dbos, status)(
                        functools.partial(func, *args, **kwargs)
                    )
                else:
                    output: R = dbos._sys_db.await_workflow_result(
                        status["workflow_uuid"],
                        polling_interval=DEFAULT_POLLING_INTERVAL,
                    )
                    return output
            except Exception as e:
                dbos.logger.error(
                    f"Exception encountered in asynchronous workflow:", exc_info=e
                )
                raise
            finally:
                if owned:
                    dbos._active_workflows_set.release(status["workflow_uuid"])


async def _execute_workflow_async(
    dbos: "DBOS",
    status: WorkflowStatusInternal,
    func: "Callable[P, Coroutine[Any, Any, R]]",
    ctx: DBOSContext,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> R:
    attributes: TracedAttributes = {
        "name": get_dbos_func_name(func),
        "operationType": OperationType.WORKFLOW.value,
    }
    fi = get_func_info(func)
    with EnterDBOSWorkflow(attributes, ctx):
        rr: Optional[str] = check_required_roles(func, fi)
        with DBOSAssumeRole(rr):
            owned = dbos._active_workflows_set.acquire(status["workflow_uuid"])
            try:
                if owned:
                    result = Pending[R](functools.partial(func, *args, **kwargs)).then(
                        _get_wf_invoke_func(dbos, status)
                    )
                    return await result()
                else:

                    def fn() -> Any:
                        return dbos._sys_db.await_workflow_result(
                            status["workflow_uuid"],
                            polling_interval=DEFAULT_POLLING_INTERVAL,
                        )

                    return await asyncio.to_thread(fn)
            except Exception as e:
                dbos.logger.error(
                    f"Exception encountered in asynchronous workflow:", exc_info=e
                )
                raise
            finally:
                if owned:
                    dbos._active_workflows_set.release(status["workflow_uuid"])


def execute_workflow_by_id(
    dbos: "DBOS", workflow_id: str, is_recovery: bool, is_dequeue: bool
) -> "WorkflowHandle[Any]":
    status = dbos._sys_db.get_workflow_status(workflow_id)
    if not status:
        raise DBOSRecoveryError(workflow_id, "Workflow status not found")
    try:
        inputs: WorkflowInputs = deserialize_args(
            status["inputs"], status["serialization"], dbos._serializer
        )
    except Exception as deser_error:
        # Mark workflow as ERROR immediately instead of leaving it PENDING for infinite retry
        try:
            error_str = serialize_exception(
                deser_error, status["serialization"], dbos._serializer
            )[0]
        except Exception:
            # Fallback: create a simple error we know can be serialized
            fallback = Exception(f"{type(deser_error).__name__}: {deser_error}")
            error_str = serialize_exception(
                fallback, status["serialization"], dbos._serializer
            )[0]
        dbos._sys_db.update_workflow_outcome(
            workflow_id,
            WorkflowStatusString.ERROR.value,
            error=error_str,
        )
        raise
    wf_func = dbos._registry.workflow_info_map.get(status["name"], None)
    if not wf_func:
        raise DBOSWorkflowFunctionNotFoundError(
            workflow_id,
            f"{status['name']} is not a registered workflow function",
        )
    fi = get_func_info(wf_func)
    if fi is None:
        raise DBOSWorkflowFunctionNotFoundError(
            "<NONE>",
            f"{wf_func.__name__} is not a registered workflow function",
        )
    # Run argument validation if configured on the workflow
    if fi.validate_args is not None and inputs is not None:
        try:
            validated_args, validated_kwargs = fi.validate_args(
                inputs["args"], inputs["kwargs"]
            )
            inputs = {"args": validated_args, "kwargs": validated_kwargs}
        except Exception as val_error:
            try:
                error_str = serialize_exception(
                    val_error, status["serialization"], dbos._serializer
                )[0]
            except Exception:
                # Fallback: create a simple error we know can be serialized
                fallback = Exception(f"{type(val_error).__name__}: {val_error}")
                error_str = serialize_exception(
                    fallback, status["serialization"], dbos._serializer
                )[0]
            dbos._sys_db.update_workflow_outcome(
                workflow_id,
                WorkflowStatusString.ERROR.value,
                error=error_str,
            )
            raise
    # Restore authentication context from the saved workflow status
    recovered_user = status.get("authenticated_user")
    recovered_roles_str = status.get("authenticated_roles")
    recovered_roles = json.loads(recovered_roles_str) if recovered_roles_str else None

    with DBOSContextSetAuth(recovered_user, recovered_roles):
        # If this function belongs to a configured class, add that class instance as its first argument
        if status["config_name"] is not None:
            config_name = status["config_name"]
            class_name = status["class_name"]
            iname = f"{class_name}/{config_name}"
            if iname not in dbos._registry.instance_info_map:
                raise DBOSWorkflowFunctionNotFoundError(
                    workflow_id,
                    f"configured class instance '{iname}' is not registered",
                )
            class_instance = dbos._registry.instance_info_map[iname]
            inputs["args"] = (class_instance,) + inputs["args"]
        # If this function is a class method, add that class object as its first argument
        elif status["class_name"] is not None:
            class_name = status["class_name"]
            if class_name not in dbos._registry.class_info_map:
                raise DBOSWorkflowFunctionNotFoundError(
                    workflow_id,
                    f"class '{class_name}' is not registered",
                )
            class_object = dbos._registry.class_info_map[class_name]
            if fi.func_type != DBOSFuncType.Static:
                inputs["args"] = (class_object,) + inputs["args"]

        with SetWorkflowID(workflow_id):
            if inspect.iscoroutinefunction(wf_func):
                ctx = get_local_dbos_context()
                parent_ctx_copy = copy.copy(ctx)
                child_ctx = DBOSContext.create_start_workflow_child(ctx)
                async_handle = dbos._background_event_loop.submit_coroutine(
                    start_workflow_async(
                        dbos,
                        parent_ctx_copy,
                        child_ctx,
                        wf_func,
                        inputs["args"],
                        inputs["kwargs"],
                        queue_name=status["queue_name"],
                        execute_workflow=True,
                        is_recovery_request=is_recovery,
                        is_dequeued_request=is_dequeue,
                    )
                )
                return WorkflowHandlePolling(async_handle.get_workflow_id(), dbos)
            return start_workflow(
                dbos,
                wf_func,
                inputs["args"],
                inputs["kwargs"],
                queue_name=status["queue_name"],
                execute_workflow=True,
                is_recovery=is_recovery,
                is_dequeued=is_dequeue,
            )


def start_workflow(
    dbos: "DBOS",
    func: "Callable[P, Union[R, Coroutine[Any, Any, R]]]",
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    queue_name: Optional[str] = None,
    execute_workflow: bool = True,
    is_recovery: bool = False,
    is_dequeued: bool = False,
) -> "WorkflowHandle[R]":

    # If the function has a class, add the class object as its first argument
    fself: Optional[object] = None
    if hasattr(func, "__self__"):
        fself = func.__self__
    if fself is not None:
        args = (fself,) + args

    fi = get_func_info(func)
    if fi is None:
        raise DBOSWorkflowFunctionNotFoundError(
            "<NONE>",
            f"{func.__name__} is not a registered workflow function",
        )
    serialization_type = fi.serialization_type
    if serialization_type is None:
        serialization_type = WorkflowSerializationFormat.DEFAULT

    func = cast("Workflow[P, R]", func.__orig_func)  # type: ignore

    inputs: WorkflowInputs = {
        "args": args,
        "kwargs": kwargs,
    }

    local_ctx = get_local_dbos_context()
    workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
        local_ctx, queue_name
    )
    workflow_timeout_ms = (
        local_ctx.workflow_timeout_ms if local_ctx is not None else None
    )
    enqueue_options = EnqueueOptionsInternal(
        deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
        priority=local_ctx.priority if local_ctx is not None else None,
        app_version=local_ctx.app_version if local_ctx is not None else None,
        queue_partition_key=(
            local_ctx.queue_partition_key if local_ctx is not None else None
        ),
    )
    new_wf_ctx = DBOSContext.create_start_workflow_child(local_ctx)
    new_child_workflow_id = new_wf_ctx.id_assigned_for_next_workflow

    if new_wf_ctx.has_parent():
        recorded_result = dbos._sys_db.check_operation_execution(
            new_wf_ctx.parent_workflow_id,
            new_wf_ctx.parent_workflow_fid,
            get_dbos_func_name(func),
        )
        if recorded_result and recorded_result["error"]:
            e: Exception = deserialize_exception(
                recorded_result["error"],
                recorded_result["serialization"],
                dbos._sys_db.serializer,
            )
            raise e
        elif recorded_result and recorded_result["child_workflow_id"]:
            return WorkflowHandlePolling(recorded_result["child_workflow_id"], dbos)

    status, should_execute = _init_workflow(
        dbos,
        new_wf_ctx,
        inputs=inputs,
        wf_name=get_dbos_func_name(func),
        class_name=get_dbos_class_name(fi, func, args),
        config_name=get_config_name(fi, func, args),
        queue=queue_name,
        workflow_timeout_ms=workflow_timeout_ms,
        workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
        max_recovery_attempts=fi.max_recovery_attempts,
        enqueue_options=enqueue_options,
        is_recovery_request=is_recovery,
        is_dequeued_request=is_dequeued,
        serialization_type=serialization_type,
    )

    if status["serialization"] == DBOSPortableJSON.name():
        serialization_type = WorkflowSerializationFormat.PORTABLE
    new_wf_ctx.serialization_type = serialization_type

    wf_status = status["status"]
    if new_wf_ctx.has_parent():
        dbos._sys_db.record_child_workflow(
            new_wf_ctx.parent_workflow_id,
            new_child_workflow_id,
            new_wf_ctx.parent_workflow_fid,
            get_dbos_func_name(func),
        )

    if (
        not execute_workflow
        or not should_execute
        or wf_status == WorkflowStatusString.ERROR.value
        or wf_status == WorkflowStatusString.SUCCESS.value
    ):
        return WorkflowHandlePolling(new_child_workflow_id, dbos)

    future = dbos._executor.submit(
        cast(Callable[..., R], _execute_workflow_wthread),
        dbos,
        status,
        func,
        new_wf_ctx,
        args,
        kwargs,
    )
    return WorkflowHandleFuture(new_child_workflow_id, future, dbos)


async def start_workflow_async(
    dbos: "DBOS",
    local_ctx: Optional[DBOSContext],
    new_wf_ctx: DBOSContext,
    func: "Callable[P, Coroutine[Any, Any, R]]",
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    queue_name: Optional[str] = None,
    execute_workflow: bool = True,
    is_recovery_request: bool = False,
    is_dequeued_request: bool = False,
) -> "WorkflowHandleAsync[R]":
    # If the function has a class, add the class object as its first argument
    fself: Optional[object] = None
    if hasattr(func, "__self__"):
        fself = func.__self__
    if fself is not None:
        args = (fself,) + args

    fi = get_func_info(func)
    if fi is None:
        raise DBOSWorkflowFunctionNotFoundError(
            "<NONE>",
            f"{func.__name__} is not a registered workflow function",
        )
    serialization_type = fi.serialization_type
    if serialization_type is None:
        serialization_type = WorkflowSerializationFormat.DEFAULT

    func = cast("Workflow[P, R]", func.__orig_func)  # type: ignore

    inputs: WorkflowInputs = {
        "args": args,
        "kwargs": kwargs,
    }

    workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
        local_ctx, queue_name
    )
    enqueue_options = EnqueueOptionsInternal(
        deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
        priority=local_ctx.priority if local_ctx is not None else None,
        app_version=local_ctx.app_version if local_ctx is not None else None,
        queue_partition_key=(
            local_ctx.queue_partition_key if local_ctx is not None else None
        ),
    )
    new_child_workflow_id = new_wf_ctx.id_assigned_for_next_workflow

    if new_wf_ctx.has_parent():
        recorded_result = await asyncio.to_thread(
            dbos._sys_db.check_operation_execution,
            new_wf_ctx.parent_workflow_id,
            new_wf_ctx.parent_workflow_fid,
            get_dbos_func_name(func),
        )
        if recorded_result and recorded_result["error"]:
            e: Exception = deserialize_exception(
                recorded_result["error"],
                recorded_result["serialization"],
                dbos._sys_db.serializer,
            )
            raise e
        elif recorded_result and recorded_result["child_workflow_id"]:
            return WorkflowHandleAsyncPolling(
                recorded_result["child_workflow_id"], dbos
            )

    status, should_execute = await asyncio.to_thread(
        _init_workflow,
        dbos,
        new_wf_ctx,
        inputs=inputs,
        wf_name=get_dbos_func_name(func),
        class_name=get_dbos_class_name(fi, func, args),
        config_name=get_config_name(fi, func, args),
        queue=queue_name,
        workflow_timeout_ms=workflow_timeout_ms,
        workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
        max_recovery_attempts=fi.max_recovery_attempts,
        enqueue_options=enqueue_options,
        is_recovery_request=is_recovery_request,
        is_dequeued_request=is_dequeued_request,
        serialization_type=serialization_type,
    )

    if status["serialization"] == DBOSPortableJSON.name():
        serialization_type = WorkflowSerializationFormat.PORTABLE
    new_wf_ctx.serialization_type = serialization_type

    if new_wf_ctx.has_parent():
        await asyncio.to_thread(
            dbos._sys_db.record_child_workflow,
            new_wf_ctx.parent_workflow_id,
            new_child_workflow_id,
            new_wf_ctx.parent_workflow_fid,
            get_dbos_func_name(func),
        )

    wf_status = status["status"]

    if (
        not execute_workflow
        or not should_execute
        or wf_status == WorkflowStatusString.ERROR.value
        or wf_status == WorkflowStatusString.SUCCESS.value
    ):
        return WorkflowHandleAsyncPolling(new_child_workflow_id, dbos)

    coro = _execute_workflow_async(dbos, status, func, new_wf_ctx, args, kwargs)
    # Shield the workflow task from cancellation
    task = asyncio.shield(asyncio.create_task(coro))
    return WorkflowHandleAsyncTask(new_child_workflow_id, task, dbos)


if sys.version_info < (3, 12):

    def _mark_coroutine(func: Callable[P, R]) -> Callable[P, R]:
        @wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> R:
            return await func(*args, **kwargs)  # type: ignore

        return async_wrapper  # type: ignore

else:

    def _mark_coroutine(func: Callable[P, R]) -> Callable[P, R]:
        inspect.markcoroutinefunction(func)

        @wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> R:
            return await func(*args, **kwargs)  # type: ignore

        return async_wrapper  # type: ignore


def workflow_wrapper(
    dbosreg: "DBOSRegistry",
    func: Callable[P, R],
    max_recovery_attempts: Optional[int] = DEFAULT_MAX_RECOVERY_ATTEMPTS,
    *,
    serialization_type: WorkflowSerializationFormat = WorkflowSerializationFormat.DEFAULT,
    validate_args: Optional["ValidateArgsCallable"] = None,
) -> Callable[P, R]:
    func.__orig_func = func  # type: ignore

    fi = get_or_create_func_info(func)
    fi.max_recovery_attempts = max_recovery_attempts
    fi.serialization_type = serialization_type
    fi.validate_args = validate_args

    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> R:
        fi = get_func_info(func)
        assert fi is not None
        if dbosreg.dbos is None:
            raise DBOSException(
                f"Function {get_dbos_func_name(func)} invoked before DBOS initialized"
            )
        dbos = dbosreg.dbos

        rr: Optional[str] = check_required_roles(func, fi)
        attributes: TracedAttributes = {
            "name": get_dbos_func_name(func),
            "operationType": OperationType.WORKFLOW.value,
        }
        inputs: WorkflowInputs = {
            "args": args,
            "kwargs": kwargs,
        }
        cctx = get_local_dbos_context()
        newwfctx = DBOSContext.create_start_workflow_child(cctx)
        resctx: Optional[DBOSContext] = None
        if cctx is not None and cctx.is_workflow():
            resctx = cctx.snapshot_step_ctx(reserve_sleep_id=False)
        workflow_timeout_ms, workflow_deadline_epoch_ms = _get_timeout_deadline(
            cctx, queue=None
        )

        wfOutcome = Outcome[R].make(functools.partial(func, *args, **kwargs))

        workflow_id = None

        def init_wf() -> Callable[[Callable[[], R]], R]:

            def recorded_result(
                c_wfid: str, dbos: "DBOS"
            ) -> Callable[[Callable[[], R]], R]:
                def recorded_result_inner(func: Callable[[], R]) -> R:
                    r: R = dbos._sys_db.await_workflow_result(
                        c_wfid, polling_interval=DEFAULT_POLLING_INTERVAL
                    )
                    return r

                return recorded_result_inner

            nonlocal workflow_id
            workflow_id = newwfctx.workflow_id

            if newwfctx.has_parent():
                r = dbos._sys_db.check_operation_execution(
                    newwfctx.parent_workflow_id,
                    newwfctx.parent_workflow_fid,
                    get_dbos_func_name(func),
                )
                if r and r["error"]:
                    e: Exception = deserialize_exception(
                        r["error"], r["serialization"], dbos._sys_db.serializer
                    )
                    raise e
                elif r and r["child_workflow_id"]:
                    return recorded_result(r["child_workflow_id"], dbos)

            status, should_execute = _init_workflow(
                dbos,
                newwfctx,
                inputs=inputs,
                wf_name=get_dbos_func_name(func),
                class_name=get_dbos_class_name(fi, func, args),
                config_name=get_config_name(fi, func, args),
                queue=None,
                workflow_timeout_ms=workflow_timeout_ms,
                workflow_deadline_epoch_ms=workflow_deadline_epoch_ms,
                max_recovery_attempts=max_recovery_attempts,
                enqueue_options=None,
                is_recovery_request=False,
                is_dequeued_request=False,
                serialization_type=fi.serialization_type,
            )

            def get_recorded_result(_func: Callable[[], R]) -> R:
                return cast(
                    R,
                    dbos._sys_db.await_workflow_result(
                        status["workflow_uuid"],
                        polling_interval=DEFAULT_POLLING_INTERVAL,
                    ),
                )

            # TODO: maybe modify the parameters if they've been changed by `_init_workflow`
            dbos.logger.debug(
                f"Running workflow, id: {newwfctx.workflow_id}, name: {get_dbos_func_name(func)}"
            )

            if newwfctx.has_parent():
                dbos._sys_db.record_child_workflow(
                    newwfctx.parent_workflow_id,
                    newwfctx.workflow_id,
                    newwfctx.parent_workflow_fid,
                    get_dbos_func_name(func),
                )

            if should_execute:
                return _get_wf_invoke_func(dbos, status)
            else:
                dbos.logger.debug(
                    f"Workflow {status['workflow_uuid']} already run with status {status['status']}"
                )
                return get_recorded_result

        def record_get_result(func: Callable[[], R]) -> R:
            """
            If a child workflow is invoked synchronously, this records the implicit "getResult" where the
            parent retrieves the child's output. It executes in the CALLER'S context, not the workflow's.
            """
            try:
                r = func()
            except Exception as e:
                serialized_e, serialization = serialize_exception(
                    e, None, dbos._serializer
                )
                assert workflow_id is not None
                dbos._sys_db.record_get_result(
                    workflow_id, None, serialized_e, serialization, resctx
                )
                raise
            serialized_r, serialization = serialize_value(r, None, dbos._serializer)
            assert workflow_id is not None
            dbos._sys_db.record_get_result(
                workflow_id, serialized_r, None, serialization, resctx
            )
            return r

        outcome = (
            wfOutcome.wrap(init_wf, dbos=dbos)
            .also(DBOSAssumeRole(rr))
            .also(EnterDBOSWorkflow(attributes, newwfctx))
            .then(record_get_result, dbos=dbos)
        )
        return outcome()  # type: ignore

    return _mark_coroutine(wrapper) if inspect.iscoroutinefunction(func) else wrapper


def decorate_workflow(
    reg: "DBOSRegistry",
    name: Optional[str],
    max_recovery_attempts: Optional[int],
    *,
    serialization_type: Optional[WorkflowSerializationFormat] = None,
    validate_args: Optional["ValidateArgsCallable"] = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    if serialization_type is None:
        serialization_type = WorkflowSerializationFormat.DEFAULT

    def _workflow_decorator(func: Callable[P, R]) -> Callable[P, R]:
        resolved_validate_args = validate_args
        # If pydantic_args_validator sentinel is passed, build a real validator
        # from the function's type hints at decoration time
        if resolved_validate_args is not None:
            from ._validation import pydantic_args_validator

            if resolved_validate_args is pydantic_args_validator:
                from ._validation import make_pydantic_args_validator

                resolved_validate_args = make_pydantic_args_validator(func)
        wrapped_func = workflow_wrapper(
            reg,
            func,
            max_recovery_attempts,
            serialization_type=serialization_type,
            validate_args=resolved_validate_args,
        )
        func_name = name if name is not None else func.__qualname__
        set_dbos_func_name(func, func_name)
        set_dbos_func_name(wrapped_func, func_name)
        reg.register_wf_function(func_name, wrapped_func, "workflow")
        return wrapped_func

    return _workflow_decorator


def decorate_transaction(
    dbosreg: "DBOSRegistry", name: Optional[str], isolation_level: "IsolationLevel"
) -> Callable[[F], F]:
    def decorator(func: F) -> F:

        transaction_name = name if name is not None else func.__qualname__

        def invoke_tx(*args: Any, **kwargs: Any) -> Any:
            if dbosreg.dbos is None:
                raise DBOSException(
                    f"Function {transaction_name} invoked before DBOS initialized"
                )
            dbos = dbosreg.dbos

            try:
                asyncio.get_running_loop()
            except RuntimeError:
                pass
            else:
                dbosreg.dbos.logger.warning(
                    f"Transaction {transaction_name} was called while an event loop is running. Invoke transactions from an async context using asyncio.to_thread to avoid blocking the event loop."
                )

            assert (
                dbos._app_db
            ), "Transactions can only be used if DBOS is configured with an application_database_url"
            with dbos._app_db.sessionmaker() as session:
                attributes: TracedAttributes = {
                    "name": transaction_name,
                    "operationType": OperationType.TRANSACTION.value,
                }
                with EnterDBOSTransaction(session, attributes=attributes):
                    ctx = assert_current_dbos_context()
                    # Check if the step record for this transaction exists
                    recorded_step_output = dbos._sys_db.check_operation_execution(
                        ctx.workflow_id, ctx.function_id, transaction_name
                    )
                    if recorded_step_output:
                        dbos.logger.debug(
                            f"Replaying transaction, id: {ctx.function_id}, name: {attributes['name']}"
                        )
                        if recorded_step_output["error"]:
                            step_error: Exception = deserialize_exception(
                                recorded_step_output["error"],
                                recorded_step_output["serialization"],
                                dbos._serializer,
                            )
                            raise step_error
                        elif recorded_step_output["output"]:
                            return deserialize_value(
                                recorded_step_output["output"],
                                recorded_step_output["child_workflow_id"],
                                dbos._serializer,
                            )
                        else:
                            raise Exception("Output and error are both None")

                    txn_output: TransactionResultInternal = {
                        "workflow_uuid": ctx.workflow_id,
                        "function_id": ctx.function_id,
                        "output": None,
                        "error": None,
                        "serialization": None,
                        "txn_snapshot": "",  # TODO: add actual snapshot
                        "executor_id": None,
                        "txn_id": None,
                        "function_name": transaction_name,
                    }
                    step_output: OperationResultInternal = {
                        "workflow_uuid": ctx.workflow_id,
                        "function_id": ctx.function_id,
                        "function_name": transaction_name,
                        "output": None,
                        "error": None,
                        "serialization": None,
                        "started_at_epoch_ms": int(time.time() * 1000),
                    }
                    retry_wait_seconds = 0.001
                    backoff_factor = 1.5
                    max_retry_wait_seconds = 2.0
                    while True:
                        has_recorded_error = False
                        txn_error: Optional[Exception] = None
                        try:
                            with session.begin():
                                # This must be the first statement in the transaction!
                                session.connection(
                                    execution_options={
                                        "isolation_level": isolation_level
                                    }
                                )
                                # Check recorded output for OAOO
                                recorded_output = (
                                    ApplicationDatabase.check_transaction_execution(
                                        session,
                                        ctx.workflow_id,
                                        ctx.function_id,
                                        transaction_name,
                                    )
                                )
                                if recorded_output:
                                    dbos.logger.debug(
                                        f"Replaying transaction, id: {ctx.function_id}, name: {attributes['name']}"
                                    )
                                    if recorded_output["error"]:
                                        deserialized_error: Exception = (
                                            deserialize_exception(
                                                recorded_output["error"],
                                                recorded_output["serialization"],
                                                dbos._serializer,
                                            )
                                        )
                                        has_recorded_error = True
                                        step_output["error"] = recorded_output["error"]
                                        step_output["serialization"] = recorded_output[
                                            "serialization"
                                        ]
                                        dbos._sys_db.record_operation_result(
                                            step_output
                                        )
                                        raise deserialized_error
                                    elif recorded_output["output"]:
                                        step_output["output"] = recorded_output[
                                            "output"
                                        ]
                                        step_output["serialization"] = recorded_output[
                                            "serialization"
                                        ]
                                        dbos._sys_db.record_operation_result(
                                            step_output
                                        )
                                        return deserialize_value(
                                            recorded_output["output"],
                                            recorded_output["serialization"],
                                            dbos._serializer,
                                        )
                                    else:
                                        raise Exception(
                                            "Output and error are both None"
                                        )
                                else:
                                    dbos.logger.debug(
                                        f"Running transaction, id: {ctx.function_id}, name: {attributes['name']}"
                                    )

                                output = func(*args, **kwargs)
                                serialized_r, serialization = serialize_value(
                                    output, None, dbos._serializer
                                )
                                txn_output["output"] = serialized_r
                                txn_output["serialization"] = serialization
                                assert (
                                    ctx.sql_session is not None
                                ), "Cannot find a database connection"
                                dbos._app_db.record_transaction_output(
                                    ctx.sql_session, txn_output
                                )
                                break
                        except DBAPIError as dbapi_error:
                            if retriable_postgres_exception(
                                dbapi_error
                            ) or dbos._app_db._is_serialization_error(dbapi_error):
                                # Retry on serialization failure
                                span = ctx.get_current_dbos_span()
                                if span:
                                    span.add_event(
                                        "Transaction Failure",
                                        {"retry_wait_seconds": retry_wait_seconds},
                                    )
                                time.sleep(retry_wait_seconds)
                                retry_wait_seconds = min(
                                    retry_wait_seconds * backoff_factor,
                                    max_retry_wait_seconds,
                                )
                                continue
                            txn_error = dbapi_error
                            raise
                        except InvalidRequestError as invalid_request_error:
                            dbos.logger.error(
                                f"InvalidRequestError in transaction {transaction_name} \033[1m Hint: Do not call commit() or rollback() within a DBOS transaction.\033[0m"
                            )
                            txn_error = invalid_request_error
                            raise
                        except DBOSUnexpectedStepError:
                            raise
                        except Exception as error:
                            txn_error = error
                            raise
                        finally:
                            # Don't record the error if it was already recorded
                            if txn_error and not has_recorded_error:
                                serialized_e, serialization = serialize_exception(
                                    txn_error, None, dbos._serializer
                                )
                                step_output["error"] = txn_output["error"] = (
                                    serialized_e
                                )
                                step_output["serialization"] = txn_output[
                                    "serialization"
                                ] = serialization
                                dbos._app_db.record_transaction_error(txn_output)
                                dbos._sys_db.record_operation_result(step_output)
            serialized_r, serialization = serialize_value(
                output, None, dbos._serializer
            )
            step_output["output"] = serialized_r
            step_output["serialization"] = serialization
            dbos._sys_db.record_operation_result(step_output)
            return output

        if inspect.iscoroutinefunction(func):
            raise DBOSException(
                f"Function {transaction_name} is a coroutine function, but DBOS.transaction does not support coroutine functions"
            )

        fi = get_or_create_func_info(func)

        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            rr: Optional[str] = check_required_roles(func, fi)
            # Entering transaction is allowed:
            #  In a workflow (that is not in a step already)
            #  Not in a workflow (we will start the single op workflow)
            ctx = get_local_dbos_context()
            if ctx and ctx.is_within_workflow():
                assert (
                    ctx.is_workflow()
                ), "Transactions must be called from within workflows"
                with DBOSAssumeRole(rr):
                    return invoke_tx(*args, **kwargs)
            else:
                tempwf = dbosreg.workflow_info_map.get("<temp>." + transaction_name)
                assert tempwf
                return tempwf(*args, **kwargs)

        set_dbos_func_name(func, transaction_name)
        set_dbos_func_name(wrapper, transaction_name)

        def temp_wf(*args: Any, **kwargs: Any) -> Any:
            return wrapper(*args, **kwargs)

        wrapped_wf = workflow_wrapper(dbosreg, temp_wf)
        set_dbos_func_name(temp_wf, "<temp>." + transaction_name)
        set_dbos_func_name(wrapped_wf, "<temp>." + transaction_name)
        set_temp_workflow_type(temp_wf, "transaction")
        dbosreg.register_wf_function(
            get_dbos_func_name(temp_wf), wrapped_wf, "transaction"
        )
        wrapper.__orig_func = temp_wf  # type: ignore
        set_func_info(wrapped_wf, get_or_create_func_info(func))
        set_func_info(temp_wf, get_or_create_func_info(func))

        return cast(F, wrapper)

    return decorator


def invoke_step(
    dbos: "DBOS",
    step_ctx: DBOSContext,
    func: Callable[P, Coroutine[Any, Any, R]] | Callable[P, R],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    *,
    step_name: str,
    retries_allowed: bool,
    interval_seconds: float,
    max_attempts: int,
    backoff_rate: float,
) -> R | Coroutine[Any, Any, R]:
    attributes: TracedAttributes = {
        "name": step_name,
        "operationType": OperationType.STEP.value,
    }

    step_start_time = int(time.time() * 1000)
    attempts = max_attempts if retries_allowed else 1
    max_retry_interval_seconds: float = 3600  # 1 Hour

    def on_exception(attempt: int, error: BaseException) -> float:
        dbos.logger.warning(
            f"Step being automatically retried (attempt {attempt + 1} of {attempts})",
            exc_info=error,
        )
        ctx = assert_current_dbos_context()
        span = ctx.get_current_dbos_span()
        if span:
            span.add_event(
                f"Step attempt {attempt} failed",
                {
                    "error": str(error),
                    "retryIntervalSeconds": interval_seconds,
                },
            )
        return min(
            interval_seconds * (backoff_rate**attempt),
            max_retry_interval_seconds,
        )

    def record_step_result(func: Callable[[], R]) -> R:
        ctx = assert_current_dbos_context()
        step_output: OperationResultInternal = {
            "workflow_uuid": ctx.workflow_id,
            "function_id": ctx.function_id,
            "function_name": step_name,
            "output": None,
            "error": None,
            "serialization": None,
            "started_at_epoch_ms": step_start_time,
        }

        try:
            output = func()
        except Exception as error:
            serialized_e, serialization = serialize_exception(
                error, None, dbos._serializer
            )
            step_output["error"] = serialized_e
            step_output["serialization"] = serialization
            dbos._sys_db.record_operation_result(step_output)
            raise
        serialized_r, serialization = serialize_value(output, None, dbos._serializer)
        step_output["output"] = serialized_r
        step_output["serialization"] = serialization
        dbos._sys_db.record_operation_result(step_output)
        return output

    def check_existing_result() -> Union[NoResult, R]:
        ctx = assert_current_dbos_context()
        recorded_output = dbos._sys_db.check_operation_execution(
            ctx.workflow_id, ctx.function_id, step_name
        )
        if recorded_output:
            dbos.logger.debug(
                f"Replaying step, id: {ctx.function_id}, name: {attributes['name']}"
            )
            if recorded_output["error"] is not None:
                deserialized_error: Exception = deserialize_exception(
                    recorded_output["error"],
                    recorded_output["serialization"],
                    dbos._serializer,
                )
                raise deserialized_error
            elif recorded_output["output"] is not None:
                return cast(
                    R,
                    deserialize_value(
                        recorded_output["output"],
                        recorded_output["serialization"],
                        dbos._serializer,
                    ),
                )
            else:
                raise Exception("Output and error are both None")
        else:
            dbos.logger.debug(
                f"Running step, id: {ctx.function_id}, name: {attributes['name']}"
            )
            return NoResult()

    stepOutcome = Outcome[R].make(functools.partial(func, *args, **kwargs))
    if retries_allowed:
        stepOutcome = stepOutcome.retry(
            max_attempts,
            on_exception,
            lambda i, e: DBOSMaxStepRetriesExceeded(step_name, i, e),
        )

    outcome = (
        stepOutcome.then(record_step_result)
        .intercept(check_existing_result, dbos=dbos)
        .also(EnterDBOSStepCtx(attributes, step_ctx))
    )
    return outcome()


def run_step(
    dbos: "DBOS",
    func: Callable[P, Coroutine[Any, Any, R]] | Callable[P, R],
    options: StepOptions,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> R:
    options = normalize_step_options(options)
    # If the step is called from a workflow, run it as a step.
    # Otherwise, run it as a normal function.
    ctx = get_local_dbos_context()
    if ctx and ctx.is_workflow():
        outcome = invoke_step(
            dbos,
            ctx.snapshot_step_ctx(),
            func,
            args,
            kwargs,
            step_name=options["name"] if options["name"] else func.__qualname__,
            retries_allowed=options["retries_allowed"],
            interval_seconds=options["interval_seconds"],
            max_attempts=options["max_attempts"],
            backoff_rate=options["backoff_rate"],
        )
        if inspect.iscoroutinefunction(func):
            return dbos._background_event_loop.submit_coroutine(
                cast(Coroutine[Any, Any, R], outcome)
            )
        else:
            return cast(R, outcome)
    else:
        if inspect.iscoroutinefunction(func):

            async def runfunc() -> R:
                return await cast(Callable[P, Coroutine[Any, Any, R]], func)(
                    *args, **kwargs
                )

            return dbos._background_event_loop.submit_coroutine(runfunc())
        else:
            return cast(Callable[P, R], func)(*args, **kwargs)


async def run_step_async(
    dbos: "DBOS",
    step_ctx: Optional[DBOSContext],
    func: Callable[P, Coroutine[Any, Any, R]] | Callable[P, R],
    options: StepOptions,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> R:
    # If the step is called from a workflow, run it as a step.
    # Otherwise, run it as a normal function.
    options = normalize_step_options(options)
    if step_ctx and step_ctx.is_workflow():
        outcome = invoke_step(
            dbos,
            step_ctx,
            func,
            args,
            kwargs,
            step_name=options["name"] if options["name"] else func.__qualname__,
            retries_allowed=options["retries_allowed"],
            interval_seconds=options["interval_seconds"],
            max_attempts=options["max_attempts"],
            backoff_rate=options["backoff_rate"],
        )
        if inspect.iscoroutinefunction(func):
            return await cast(Coroutine[Any, Any, R], outcome)
        else:
            return cast(R, outcome)
    else:
        if inspect.iscoroutinefunction(func):
            return await cast(Callable[P, Coroutine[Any, Any, R]], func)(
                *args, **kwargs
            )
        else:
            return await asyncio.to_thread(
                lambda: cast(Callable[P, R], func)(*args, **kwargs)
            )


def decorate_step(
    dbosreg: "DBOSRegistry",
    *,
    name: Optional[str],
    retries_allowed: bool,
    interval_seconds: float,
    max_attempts: int,
    backoff_rate: float,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    def decorator(func: Callable[P, R]) -> Callable[P, R]:

        step_name = name if name is not None else func.__qualname__

        fi = get_or_create_func_info(func)

        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            # If the step is called from a workflow, run it as a step.
            # Otherwise, run it as a normal function.
            ctx = get_local_dbos_context()
            if ctx and ctx.is_workflow():
                if dbosreg.dbos is None:
                    raise DBOSException(
                        f"Function {step_name} invoked before DBOS initialized"
                    )
                rr: Optional[str] = check_required_roles(func, fi)
                with DBOSAssumeRole(rr):
                    return invoke_step(
                        dbosreg.dbos,
                        ctx.snapshot_step_ctx(),
                        func,
                        args,
                        kwargs,
                        step_name=step_name,
                        retries_allowed=retries_allowed,
                        interval_seconds=interval_seconds,
                        max_attempts=max_attempts,
                        backoff_rate=backoff_rate,
                    )
            else:
                return func(*args, **kwargs)

        wrapper = (
            _mark_coroutine(wrapper) if inspect.iscoroutinefunction(func) else wrapper  # type: ignore
        )

        set_dbos_func_name(func, step_name)
        set_dbos_func_name(wrapper, step_name)

        def temp_wf_sync(*args: Any, **kwargs: Any) -> Any:
            return wrapper(*args, **kwargs)

        async def temp_wf_async(*args: Any, **kwargs: Any) -> Any:
            return await wrapper(*args, **kwargs)

        temp_wf = temp_wf_async if inspect.iscoroutinefunction(func) else temp_wf_sync
        wrapped_wf = workflow_wrapper(dbosreg, temp_wf)
        set_dbos_func_name(temp_wf, "<temp>." + step_name)
        set_dbos_func_name(wrapped_wf, "<temp>." + step_name)
        set_temp_workflow_type(temp_wf, "step")
        dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf, "step")
        wrapper.__orig_func = temp_wf  # type: ignore
        set_func_info(wrapped_wf, get_or_create_func_info(func))
        set_func_info(temp_wf, get_or_create_func_info(func))

        return cast(Callable[P, R], wrapper)

    return decorator


def send(
    dbos: "DBOS",
    cur_ctx: Optional["DBOSContext"],
    destination_id: str,
    message: Any,
    topic: Optional[str] = None,
    *,
    serialization_type: Optional[WorkflowSerializationFormat],
    idempotency_key: Optional[str] = None,
) -> None:
    if (
        serialization_type is None
        or serialization_type == WorkflowSerializationFormat.DEFAULT
    ):
        serialization_type = (
            cur_ctx.serialization_type
            if cur_ctx is not None
            else WorkflowSerializationFormat.DEFAULT
        )

    def do_send(destination_id: str, message: Any, topic: Optional[str]) -> None:
        assert cur_ctx is not None
        attributes: TracedAttributes = {
            "name": "send",
        }
        with EnterDBOSStepCtx(attributes, cur_ctx) as ctx:
            dbos._sys_db.send(
                ctx.workflow_id,
                ctx.curr_step_function_id,
                destination_id,
                message,
                topic,
                serialization_type=serialization_type,
                message_uuid=idempotency_key,
            )

    if cur_ctx and cur_ctx.is_workflow():
        return do_send(destination_id, message, topic)
    else:
        dbos._sys_db.send_direct(
            destination_id,
            message,
            topic,
            message_uuid=idempotency_key,
            serialization_type=serialization_type,
        )


def set_event(
    dbos: "DBOS",
    cur_ctx: Optional["DBOSContext"],
    key: str,
    value: Any,
    *,
    serialization_type: WorkflowSerializationFormat,
) -> None:
    if (
        serialization_type is None
        or serialization_type == WorkflowSerializationFormat.DEFAULT
    ):
        serialization_type = (
            cur_ctx.serialization_type
            if cur_ctx is not None
            else WorkflowSerializationFormat.DEFAULT
        )

    if cur_ctx is not None:
        if cur_ctx.is_workflow():
            # If called from a workflow function, run as a step
            attributes: TracedAttributes = {
                "name": "set_event",
            }
            with EnterDBOSStepCtx(attributes, cur_ctx) as ctx:
                dbos._sys_db.set_event_from_workflow(
                    ctx.workflow_id,
                    ctx.curr_step_function_id,
                    key,
                    value,
                    serialization_type=serialization_type,
                )
        elif cur_ctx.is_step():
            dbos._sys_db.set_event_from_step(
                cur_ctx.workflow_id,
                cur_ctx.curr_step_function_id,
                key,
                value,
                serialization_type=serialization_type,
            )
        else:
            raise DBOSException(
                "set_event() must be called from within a workflow or step"
            )
    else:
        raise DBOSException("set_event() must be called from within a workflow or step")


def record_sleep(
    dbos: "DBOS", cur_ctx: Optional["DBOSContext"], seconds: float
) -> float:
    if cur_ctx is not None:
        # Must call it within a workflow
        assert cur_ctx.is_workflow(), "sleep() must be called from within a workflow"
        attributes: TracedAttributes = {
            "name": "sleep",
        }
        with EnterDBOSStepCtx(attributes, cur_ctx) as ctx:
            return dbos._sys_db.record_sleep(
                ctx.workflow_id, ctx.curr_step_function_id, seconds
            )
    else:
        # Cannot call it from outside of a workflow
        raise DBOSException("sleep() must be called from within a workflow")


def write_stream(
    dbos: "DBOS",
    step_ctx: Optional["DBOSContext"],
    key: str,
    value: Any,
    *,
    serialization_type: WorkflowSerializationFormat,
) -> None:
    if (
        serialization_type is None
        or serialization_type == WorkflowSerializationFormat.DEFAULT
    ):
        serialization_type = (
            step_ctx.serialization_type
            if step_ctx is not None
            else WorkflowSerializationFormat.DEFAULT
        )

    if step_ctx is not None:
        # Must call it within a workflow
        if step_ctx.is_workflow():
            attributes: TracedAttributes = {
                "name": "write_stream",
            }
            with EnterDBOSStepCtx(attributes, step_ctx) as ctx:
                dbos._sys_db.write_stream_from_workflow(
                    ctx.workflow_id,
                    ctx.function_id,
                    key,
                    value,
                    serialization_type=serialization_type,
                )
        elif step_ctx.is_step():
            dbos._sys_db.write_stream_from_step(
                step_ctx.workflow_id,
                step_ctx.function_id,
                key,
                value,
                serialization_type=serialization_type,
            )
        else:
            raise DBOSException(
                "write_stream() must be called from within a workflow or step"
            )
    else:
        # Cannot call it from outside of a workflow
        raise DBOSException(
            "write_stream() must be called from within a workflow or step"
        )


def close_stream(dbos: "DBOS", step_ctx: Optional["DBOSContext"], key: str) -> None:
    if step_ctx is not None:
        # Must call it within a workflow
        if step_ctx.is_workflow():
            attributes: TracedAttributes = {
                "name": "close_stream",
            }
            with EnterDBOSStepCtx(attributes, step_ctx) as ctx:
                dbos._sys_db.close_stream(ctx.workflow_id, ctx.function_id, key)
        else:
            raise DBOSException("close_stream() must be called from within a workflow")
    else:
        # Cannot call it from outside of a workflow
        raise DBOSException("close_stream() must be called from within a workflow")


def _get_timeout_deadline(
    ctx: Optional[DBOSContext], queue: Optional[str]
) -> tuple[Optional[int], Optional[int]]:
    if ctx is None:
        return None, None
    # If a timeout is explicitly specified, use it over any propagated deadline
    if ctx.workflow_timeout_ms:
        if queue:
            # Queued workflows are assigned a deadline on dequeue
            return ctx.workflow_timeout_ms, None
        else:
            # Otherwise, compute the deadline immediately
            return (
                ctx.workflow_timeout_ms,
                int(time.time() * 1000) + ctx.workflow_timeout_ms,
            )
    # Otherwise, return the propagated deadline, if any
    else:
        return None, ctx.workflow_deadline_epoch_ms
