task.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

# See https://peps.python.org/pep-0563/
from __future__ import annotations

import logging
import math
from abc import ABC, abstractmethod
from collections.abc import Callable, Generator, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeVar, cast, overload

from durabletask.entities import DurableEntity, EntityInstanceId, EntityLock, EntityContext
import durabletask.internal.helpers as pbh
import durabletask.internal.orchestrator_service_pb2 as pb

T = TypeVar('T')
TInput = TypeVar('TInput')
TOutput = TypeVar('TOutput')


class OrchestrationContext(ABC):

    @property
    @abstractmethod
    def instance_id(self) -> str:
        """Get the ID of the current orchestration instance.

        The instance ID is generated and fixed when the orchestrator function
        is scheduled. It can be either auto-generated, in which case it is
        formatted as a UUID, or it can be user-specified with any format.

        Returns
        -------
        str
            The ID of the current orchestration instance.
        """
        pass

    @property
    @abstractmethod
    def parent_instance_id(self) -> str | None:
        """Get the ID of the parent orchestration instance.

        For a sub-orchestration, this is the instance ID of the orchestration
        that scheduled it. For a top-level orchestration, this is ``None``.

        Returns
        -------
        str | None
            The parent orchestration instance ID, or ``None`` if this
            orchestration was not scheduled by a parent orchestration.
        """
        pass

    @property
    @abstractmethod
    def version(self) -> str | None:
        """Get the version of the orchestration instance.

        This version is set when the orchestration is scheduled and can be used
        to determine which version of the orchestrator function is being executed.

        Returns
        -------
        str | None
            The version of the orchestration instance, or None if not set.
        """
        pass

    @property
    @abstractmethod
    def current_utc_datetime(self) -> datetime:
        """Get the current date/time as UTC.

        This date/time value is derived from the orchestration history. It
        always returns the same value at specific points in the orchestrator
        function code, making it deterministic and safe for replay.

        Returns
        -------
        datetime
            The current timestamp in a way that is safe for use by orchestrator functions
        """
        pass

    @property
    @abstractmethod
    def is_replaying(self) -> bool:
        """Get the value indicating whether the orchestrator is replaying from history.

        This property is useful when there is logic that needs to run only when
        the orchestrator function is _not_ replaying. For example, certain
        types of application logging may become too noisy when duplicated as
        part of orchestrator function replay. The orchestrator code could check
        to see whether the function is being replayed and then issue the log
        statements when this value is `false`.

        Returns
        -------
        bool
            Value indicating whether the orchestrator function is currently replaying.
        """
        pass

    @abstractmethod
    def set_custom_status(self, custom_status: Any) -> None:
        """Set the orchestration instance's custom status.

        Parameters
        ----------
        custom_status: Any
            A JSON-serializable custom status value to set.
        """
        pass

    @abstractmethod
    def create_timer(self, fire_at: datetime | timedelta) -> TimerTask:
        """Create a Timer Task to fire after at the specified deadline.

        Parameters
        ----------
        fire_at: datetime.datetime | datetime.timedelta
            The time for the timer to trigger or a time delta from now.

        Returns
        -------
        TimerTask
            A Durable Timer Task that schedules the timer to wake up the orchestrator
        """
        pass

    @overload
    def call_activity(self, activity: Activity[TInput, TOutput] | str, *,
                      input: TInput | None = ...,
                      retry_policy: RetryPolicy | None = ...,
                      tags: dict[str, str] | None = ...,
                      return_type: type[T]) -> CompletableTask[T]:
        ...

    @overload
    def call_activity(self, activity: Activity[TInput, TOutput] | str, *,
                      input: TInput | None = ...,
                      retry_policy: RetryPolicy | None = ...,
                      tags: dict[str, str] | None = ...,
                      return_type: None = ...) -> CompletableTask[TOutput]:
        ...

    @abstractmethod
    def call_activity(self, activity: Activity[TInput, TOutput] | str, *,
                      input: TInput | None = None,
                      retry_policy: RetryPolicy | None = None,
                      tags: dict[str, str] | None = None,
                      return_type: type | None = None) -> CompletableTask[Any]:
        """Schedule an activity for execution.

        Parameters
        ----------
        activity: Activity[TInput, TOutput] | str
            A reference to the activity function to call.
        input: TInput | None
            The JSON-serializable input (or None) to pass to the activity.
        retry_policy: RetryPolicy | None
            The retry policy to use for this activity call.
        tags: dict[str, str] | None
            Optional tags to associate with the activity invocation.
        return_type: type | None
            Optional type used to deserialize the activity's result. When
            provided, the result is coerced to this type (dataclasses are
            constructed from their dict payloads, types exposing a
            ``from_json()`` classmethod are reconstructed via that hook), and
            the returned task is typed as ``CompletableTask[return_type]``.
            When omitted, the return type is discovered from the activity
            function's return annotation (if a function reference is passed and
            it is annotated with a reconstructable type); otherwise the raw
            deserialized JSON is returned.

        Returns
        -------
        Task
            A Durable Task that completes when the called activity function completes or fails.
        """
        pass

    @overload
    def call_entity(self,
                    entity: EntityInstanceId,
                    operation: str,
                    input: Any = ...,
                    *,
                    return_type: type[T]) -> CompletableTask[T]:
        ...

    @overload
    def call_entity(self,
                    entity: EntityInstanceId,
                    operation: str,
                    input: Any = ...,
                    *,
                    return_type: None = ...) -> CompletableTask[Any]:
        ...

    @abstractmethod
    def call_entity(self,
                    entity: EntityInstanceId,
                    operation: str,
                    input: Any = None,
                    *,
                    return_type: type | None = None) -> CompletableTask[Any]:
        """Schedule entity function for execution.

        Parameters
        ----------
        entity: EntityInstanceId
            The ID of the entity instance to call.
        operation: str
            The name of the operation to invoke on the entity.
        input: TInput | None
            The optional JSON-serializable input to pass to the entity function.
        return_type: type | None
            Optional type used to deserialize the operation's result. When
            provided, the result is coerced to this type and the returned task
            is typed as ``CompletableTask[return_type]``; when omitted, the raw
            deserialized JSON is returned.

        Returns
        -------
        Task
            A Durable Task that completes when the called entity function completes or fails.
        """
        pass

    @abstractmethod
    def signal_entity(
            self,
            entity_id: EntityInstanceId,
            operation_name: str,
            input: Any = None,
            signal_time: datetime | None = None
    ) -> None:
        """Signal an entity function for execution.

        Parameters
        ----------
        entity_id: EntityInstanceId
            The ID of the entity instance to signal.
        operation_name: str
            The name of the operation to invoke on the entity.
        input: TInput | None
            The optional JSON-serializable input to pass to the entity function.
        signal_time: datetime | None
            The optional time at which the signal should be delivered. If None, the
            signal is delivered as soon as possible. Use this to schedule a future
            operation on the entity.
        """
        pass

    @abstractmethod
    def lock_entities(self, entities: list[EntityInstanceId]) -> CompletableTask[EntityLock]:
        """Creates a Task object that locks the specified entity instances.

        The locks will be acquired the next time the orchestrator yields.
        Best practice is to immediately yield this Task and enter the returned EntityLock.
        The lock is released when the EntityLock is exited.

        Parameters
        ----------
        entities: list[EntityInstanceId]
            The list of entity instance IDs to lock.

        Returns
        -------
        EntityLock
            A context manager object that releases the locks when exited.
        """
        pass

    @overload
    def call_sub_orchestrator(self, orchestrator: Orchestrator[TInput, TOutput] | str, *,
                              input: TInput | None = ...,
                              instance_id: str | None = ...,
                              retry_policy: RetryPolicy | None = ...,
                              version: str | None = ...,
                              return_type: type[T]) -> CompletableTask[T]:
        ...

    @overload
    def call_sub_orchestrator(self, orchestrator: Orchestrator[TInput, TOutput] | str, *,
                              input: TInput | None = ...,
                              instance_id: str | None = ...,
                              retry_policy: RetryPolicy | None = ...,
                              version: str | None = ...,
                              return_type: None = ...) -> CompletableTask[TOutput]:
        ...

    @abstractmethod
    def call_sub_orchestrator(self, orchestrator: Orchestrator[TInput, TOutput] | str, *,
                              input: TInput | None = None,
                              instance_id: str | None = None,
                              retry_policy: RetryPolicy | None = None,
                              version: str | None = None,
                              return_type: type | None = None) -> CompletableTask[Any]:
        """Schedule sub-orchestrator function for execution.

        Parameters
        ----------
        orchestrator: Orchestrator[TInput, TOutput]
            A reference to the orchestrator function to call.
        input: TInput | None
            The optional JSON-serializable input to pass to the orchestrator function.
        instance_id: str | None
            A unique ID to use for the sub-orchestration instance. If not specified, a
            random UUID will be used.
        retry_policy: RetryPolicy | None
            The retry policy to use for this sub-orchestrator call.
        return_type: type | None
            Optional type used to deserialize the sub-orchestrator's result. When
            provided, the result is coerced to this type and the returned task is
            typed as ``CompletableTask[return_type]``; when omitted, the raw
            deserialized JSON is returned.

        Returns
        -------
        Task
            A Durable Task that completes when the called sub-orchestrator completes or fails.
        """
        pass

    # TOOD: Add a timeout parameter, which allows the task to be cancelled if the event is
    # not received within the specified timeout. This requires support for task cancellation.
    @overload
    def wait_for_external_event(self, name: str, *,
                                data_type: type[T]) -> CancellableTask[T]:
        ...

    @overload
    def wait_for_external_event(self, name: str, *,
                                data_type: None = ...) -> CancellableTask[Any]:
        ...

    @abstractmethod
    def wait_for_external_event(self, name: str, *,
                                data_type: type | None = None) -> CancellableTask[Any]:
        """Wait asynchronously for an event to be raised with the name `name`.

        Parameters
        ----------
        name : str
            The event name of the event that the task is waiting for.
        data_type : type | None
            Optional type used to deserialize the event payload. When provided,
            the payload is coerced to this type and the returned task is typed
            as ``CancellableTask[data_type]``; when omitted, the raw
            deserialized JSON is returned.

        Returns
        -------
        CancellableTask[Any]
            A Durable Task that completes when the event is received.
        """
        pass

    def send_event(self, instance_id: str, event_name: str, *,
                   data: Any | None = None) -> None:
        """Send an event to another orchestration instance.

        The target orchestration can receive the event using
        :meth:`wait_for_external_event`. This is a one-way operation and does
        not wait for the target orchestration to process the event. If the
        target orchestration does not exist, the event is silently dropped.
        During replay, the Python SDK validates both the event name and target
        instance ID. This is intentionally stricter than DurableTask.Core,
        which does not validate the target instance ID.

        Parameters
        ----------
        instance_id : str
            The ID of the orchestration instance to send the event to.
        event_name : str
            The name of the event to send. Event names are case-insensitive.
        data : Any | None
            The optional serializable event payload.
        """
        raise NotImplementedError(
            "This OrchestrationContext implementation does not support "
            "send_event()."
        )

    @abstractmethod
    def continue_as_new(self, new_input: Any, *, save_events: bool = False,
                        new_version: str | None = None) -> None:
        """Continue the orchestration execution as a new instance.

        Orchestrators should return immediately after calling this method.
        Subsequent external events are no longer delivered to pending waits
        in the current execution.

        Parameters
        ----------
        new_input : Any
            The new input to use for the new orchestration instance.
        save_events : bool
            A flag indicating whether to add any unprocessed external events in the new orchestration history.
            Events already delivered to a waiting task are not saved, even if
            the orchestrator has not yielded that task.
        new_version : str | None
            An optional version to assign to the new orchestration instance.
        """
        pass

    @abstractmethod
    def new_uuid(self) -> str:
        """Create a new UUID that is safe for replay within an orchestration or operation.

        The default implementation of this method creates a name-based UUID
        using the algorithm from RFC 4122 §4.3. The name input used to generate
        this value is a combination of the orchestration instance ID, the current UTC datetime,
        and an internally managed counter.

        Returns
        -------
        str
            New UUID that is safe for replay within an orchestration or operation.
        """
        pass

    @abstractmethod
    def _exit_critical_section(self) -> None:
        pass

    def create_replay_safe_logger(self, logger: logging.Logger) -> ReplaySafeLogger:
        """Create a replay-safe logger that suppresses log messages during orchestration replay.

        The returned logger wraps the provided logger and only emits log messages when
        the orchestrator is not replaying. This prevents duplicate log messages from
        appearing as a side effect of orchestration replay.

        Parameters
        ----------
        logger : logging.Logger
            The underlying logger to wrap.

        Returns
        -------
        ReplaySafeLogger
            A logger that only emits log messages when the orchestrator is not replaying.
        """
        return ReplaySafeLogger(logger, lambda: self.is_replaying)


if TYPE_CHECKING:
    # logging.LoggerAdapter is generic in stubs but is not subscriptable
    # at runtime before Python 3.11. Use a TYPE_CHECKING alias so the
    # base class evaluates correctly at runtime.
    _LoggerAdapterBase = logging.LoggerAdapter[logging.Logger]
else:
    _LoggerAdapterBase = logging.LoggerAdapter


class ReplaySafeLogger(_LoggerAdapterBase):
    """A logger adapter that suppresses log messages during orchestration replay.

    This class extends :class:`logging.LoggerAdapter` and only emits log
    messages when the orchestrator is *not* replaying. Use this to avoid
    duplicate log entries that would otherwise appear every time the
    orchestrator replays its history.

    Obtain an instance by calling :meth:`OrchestrationContext.create_replay_safe_logger`.
    """

    def __init__(self, logger: logging.Logger, is_replaying: Callable[[], bool]) -> None:
        super().__init__(logger, {})
        self._is_replaying = is_replaying

    def isEnabledFor(self, level: int) -> bool:
        """Return whether logging is enabled for the given level.

        Returns ``False`` while the orchestrator is replaying so that callers
        can skip expensive message formatting during replay.
        """
        if self._is_replaying():
            return False
        return self.logger.isEnabledFor(level)


@dataclass(frozen=True)
class FailureDetails:
    message: str
    error_type: str
    stack_trace: str | None
    inner_failure: FailureDetails | None = None
    properties: dict[str, Any] | None = None

    def is_caused_by(self, error_type: str | type[BaseException]) -> bool:
        """Return ``True`` if this failure was caused by ``error_type``.

        This mirrors the .NET ``TaskFailureDetails.IsCausedBy<T>()`` and Java
        ``FailureDetails.isCausedBy(Class)`` helpers, letting a caller introspect
        a failure instead of comparing :attr:`error_type` strings by hand::

            try:
                yield ctx.call_activity(my_activity)
            except TaskFailedError as e:
                if e.details.is_caused_by(ValueError):
                    ...

        ``error_type`` may be either an exception type or a name string:

        * When an exception **type** is given, the check is base-type aware: it
          returns ``True`` if the failure's type is ``error_type`` or any of its
          subclasses. To stay safe, this only inspects exception subclasses that
          are already imported in the current process (it never imports a type by
          name), so a subclass whose module has not been imported yet will not be
          matched. Passing a non-exception type raises :class:`TypeError`.
        * When a **string** is given, it is compared by name only (no base-type
          awareness). The name may be fully qualified (``"module.ClassName"``) or
          unqualified (``"ClassName"``); an unqualified name matches on the class
          name alone and therefore cannot distinguish same-named types from
          different modules.

        Only this failure's own :attr:`error_type` is considered; chained/inner
        causes are not traversed.
        """
        if isinstance(error_type, str):
            return self._name_matches(error_type)
        # The annotation restricts callers to exception types, but this is a
        # public API, so validate at runtime against arbitrary objects.
        if not (isinstance(error_type, type) and issubclass(error_type, BaseException)):  # pyright: ignore[reportUnnecessaryIsInstance]
            raise TypeError(
                "is_caused_by() expects an exception type or a type-name string, "
                f"but got {error_type!r}.")

        # Base-type-aware match without importing the failure's type: walk the
        # target type and its already-loaded subclasses (visited-marked to guard
        # against diamond inheritance) and compare each candidate's qualified
        # name against this failure's error_type.
        visited: set[int] = set()
        stack: list[type] = [error_type]
        while stack:
            candidate = stack.pop()
            if id(candidate) in visited:
                continue
            visited.add(id(candidate))
            if self._type_matches(candidate):
                return True
            stack.extend(candidate.__subclasses__())
        return False

    def _type_matches(self, candidate: type) -> bool:
        if not self.error_type:
            return False
        candidate_name = pbh.get_qualified_name(candidate)
        if self.error_type == candidate_name:
            return True
        # Back-compat: a bare (non-qualified) stored name -- e.g. produced by an
        # older SDK version or a non-Python SDK -- can only be compared by the
        # unqualified class name.
        if "." not in self.error_type:
            return self.error_type == candidate_name.rsplit(".", 1)[-1]
        return False

    def _name_matches(self, name: str) -> bool:
        stored = self.error_type
        if not stored or not name:
            return False
        if stored == name:
            return True
        # When both names are fully qualified, require an exact match so that
        # same-named types from different modules are not confused. Otherwise,
        # honor the "qualified or unqualified" contract by comparing the
        # unqualified (trailing) segment.
        if "." in stored and "." in name:
            return False
        return stored.rsplit(".", 1)[-1] == name.rsplit(".", 1)[-1]


class TaskFailedError(Exception):
    """Exception type for all orchestration task failures."""

    def __init__(self, message: str, details: pb.TaskFailureDetails | Exception):
        super().__init__(message)
        if isinstance(details, Exception):
            nested_failure = getattr(details, "failure_details", None)
            details = pbh.new_failure_details(details)
            if isinstance(nested_failure, pb.TaskFailureDetails):
                details.innerFailure.CopyFrom(nested_failure)
                for key, value in nested_failure.properties.items():
                    details.properties[key].CopyFrom(value)
        self._failure_details = details
        self._details = pbh.failure_details_from_protobuf(details)

    @property
    def details(self) -> FailureDetails:
        return self._details


class NonDeterminismError(Exception):
    pass


class OrchestrationStateError(Exception):
    pass


class TaskCancelledError(Exception):
    """Exception type for cancelled orchestration tasks."""


class _TaskCompletionError(Exception):
    """Collect callback errors without requiring Python 3.11 ExceptionGroup."""

    def __init__(self, errors: list[Exception]) -> None:
        self.errors = tuple(errors)
        super().__init__(
            "Multiple task completion callbacks failed: "
            + "; ".join(f"{type(error).__name__}: {error}" for error in errors))


class Task(ABC, Generic[T]):
    """Abstract base class for asynchronous tasks in a durable orchestration."""
    _result: T
    _exception: TaskFailedError | None
    _parents: dict[int, CompositeTask[Any]]

    def __init__(self) -> None:
        super().__init__()
        self._is_complete = False
        self._exception = None
        # Dict insertion order preserves subscription order for shared races.
        self._parents = {}

    @property
    def is_complete(self) -> bool:
        """Returns True if the task has completed, False otherwise."""
        return self._is_complete

    @property
    def is_failed(self) -> bool:
        """Returns True if the task has failed, False otherwise."""
        return self._exception is not None

    @property
    def result(self) -> T:
        """Returns the result of the task (alias for :meth:`get_result`)."""
        return self.get_result()

    def get_result(self) -> T:
        """Returns the result of the task."""
        if not self._is_complete:
            raise ValueError('The task has not completed.')
        elif self._exception is not None:
            raise self._exception
        return self._result

    def get_exception(self) -> TaskFailedError:
        """Returns the exception that caused the task to fail."""
        if self._exception is None:
            raise ValueError('The task has not failed.')
        return self._exception

    def _notify_parent(self) -> None:
        parents = tuple(self._parents.values())
        self._parents.clear()
        errors: list[Exception] = []
        for parent in parents:
            try:
                parent.on_child_completed(self)
            except Exception as error:
                # A canceled result can throw without blocking other observers.
                errors.append(error)
        if len(errors) == 1:
            raise errors[0]
        if errors:
            raise _TaskCompletionError(errors) from errors[0]


class CompositeTask(Task[T]):
    """A task that is composed of other tasks."""
    _tasks: list[Task[Any]]
    _completed_tasks: int
    _failed_tasks: int

    def __init__(self, tasks: list[Task[Any]]):
        super().__init__()
        self._tasks = tasks
        self._completed_tasks = 0
        self._failed_tasks = 0
        # Track references by identity, including unhashable Task subclasses.
        self._pending_children: dict[int, tuple[Task[Any], int]] = {}
        for task in tasks:
            _, count = self._pending_children.get(id(task), (task, 0))
            self._pending_children[id(task)] = (task, count + 1)
        for task, _ in tuple(self._pending_children.values()):
            if self.is_complete:
                break
            if task.is_complete:
                self.on_child_completed(task)
            else:
                task._parents[id(self)] = self

    def _notify_parent(self) -> None:
        for task, _ in self._pending_children.values():
            task._parents.pop(id(self), None)
        self._pending_children.clear()
        super()._notify_parent()

    def get_tasks(self) -> list[Task[Any]]:
        return self._tasks

    @abstractmethod
    def on_child_completed(self, task: Task[Any]) -> None:
        pass


class WhenAllTask(CompositeTask[list[T]]):
    """A task that completes when all of its child tasks complete."""

    def __init__(self, tasks: list[Task[T]]):
        # Initialize state that on_child_completed() reads BEFORE invoking the
        # base constructor: CompositeTask.__init__ calls on_child_completed()
        # for any children that are already complete, so `_pending_exception`
        # must exist first. The base constructor also initializes
        # `_completed_tasks`/`_failed_tasks` to 0 and then accounts for
        # pre-completed children, so they must not be reset afterwards.
        self._pending_exception: TaskFailedError | None = None
        super().__init__(cast(list[Task[Any]], tasks))

    @property
    def pending_tasks(self) -> int:
        """Returns the number of tasks that have not yet completed."""
        return len(self._tasks) - self._completed_tasks

    def on_child_completed(self, task: Task[Any]) -> None:
        if self.is_complete:
            return
        pending = self._pending_children.pop(id(task), None)
        if pending is None:
            return
        _, occurrences = pending
        self._completed_tasks += occurrences
        if task.is_failed:
            self._failed_tasks += occurrences
            if self._pending_exception is None:
                # Stage the first failure but do NOT expose it via `_exception`
                # yet. Exposing it now would make `is_failed` return True while
                # `is_complete` is still False, diverging from .NET's
                # Task.WhenAll (which does not fault until all children finish).
                self._pending_exception = task.get_exception()
        if self._completed_tasks == len(self._tasks):
            # Only complete once every child task has completed. This matches the
            # semantics of .NET's Task.WhenAll: the composite task waits for all
            # children to finish and, if any failed, surfaces the first failure
            # rather than failing fast on the first error.
            self._is_complete = True
            if self._pending_exception is not None:
                self._exception = self._pending_exception
            else:
                # The order of the result MUST match the order of the tasks
                # provided to the constructor.
                self._result = [child.get_result() for child in self._tasks]
            self._notify_parent()

    def get_completed_tasks(self) -> int:
        return self._completed_tasks


class CompletableTask(Task[T]):

    def __init__(self, expected_type: type | None = None) -> None:
        super().__init__()
        self._retryable_parent: RetryableTask[Any] | None = None
        self._expected_type = expected_type

    def complete(self, result: T):
        if self._is_complete:
            raise ValueError('The task has already completed.')
        self._result = result
        self._is_complete = True
        self._notify_parent()

    def fail(self, message: str, details: Exception | pb.TaskFailureDetails):
        if self._is_complete:
            raise ValueError('The task has already completed.')
        self._exception = TaskFailedError(message, details)
        self._is_complete = True
        self._notify_parent()


class CancellableTask(CompletableTask[T]):
    """A completable task that can be cancelled before it finishes."""

    def __init__(self, expected_type: type | None = None) -> None:
        super().__init__(expected_type)
        self._is_cancelled = False
        self._cancel_handler: Callable[[], None] | None = None

    @property
    def is_cancelled(self) -> bool:
        """Returns True if the task was cancelled, False otherwise."""
        return self._is_cancelled

    def get_result(self) -> T:
        if self._is_cancelled:
            raise TaskCancelledError('The task was cancelled.')
        return super().get_result()

    def set_cancel_handler(self, cancel_handler: Callable[[], None]) -> None:
        self._cancel_handler = cancel_handler

    def cancel(self) -> bool:
        """Attempts to cancel this task.

        Returns
        -------
        bool
            True if cancellation was applied, False if the task had already completed.
        """
        if self._is_complete:
            return False

        if self._cancel_handler is not None:
            self._cancel_handler()

        self._is_cancelled = True
        self._is_complete = True
        self._notify_parent()
        return True


class RetryableTask(CompletableTask[T]):
    """A task that can be retried according to a retry policy."""

    def __init__(self, retry_policy: RetryPolicy, action: pb.OrchestratorAction,
                 start_time: datetime, is_sub_orch: bool,
                 expected_type: type | None = None) -> None:
        super().__init__(expected_type)
        self._action = action
        self._retry_policy = retry_policy
        self._attempt_count = 1
        self._start_time = start_time
        self._is_sub_orch = is_sub_orch

    def increment_attempt_count(self) -> None:
        self._attempt_count += 1

    def compute_next_delay(self) -> timedelta | None:
        if self._attempt_count >= self._retry_policy.max_number_of_attempts:
            return None

        retry_expiration: datetime = datetime.max
        if self._retry_policy.retry_timeout is not None and self._retry_policy.retry_timeout != datetime.max:
            retry_expiration = self._start_time + self._retry_policy.retry_timeout

        if self._retry_policy.backoff_coefficient is None:
            backoff_coefficient = 1.0
        else:
            backoff_coefficient = self._retry_policy.backoff_coefficient

        if datetime.now(tz=timezone.utc).replace(tzinfo=None) < retry_expiration:
            next_delay_f = math.pow(backoff_coefficient, self._attempt_count - 1) * self._retry_policy.first_retry_interval.total_seconds()

            if self._retry_policy.max_retry_interval is not None:
                next_delay_f = min(next_delay_f, self._retry_policy.max_retry_interval.total_seconds())

            return timedelta(seconds=next_delay_f)

        return None


class TimerTask(CancellableTask[None]):
    def __init__(self, final_fire_at: datetime | None = None,
                 maximum_timer_interval: timedelta | None = None):
        super().__init__()
        self._final_fire_at = final_fire_at
        self._maximum_timer_interval = maximum_timer_interval

    def set_retryable_parent(self, retryable_task: RetryableTask[Any]) -> None:
        self._retryable_parent = retryable_task

    def _handle_timer_fired(self, current_utc_datetime: datetime) -> datetime | None:
        if (self._final_fire_at is not None
                and self._maximum_timer_interval is not None
                and current_utc_datetime < self._final_fire_at):
            return self._get_next_fire_at(current_utc_datetime)
        super().complete(None)
        return None

    def _get_next_fire_at(self, current_utc_datetime: datetime) -> datetime:
        # _handle_timer_fired guards both attributes before calling this method.
        assert self._final_fire_at is not None
        assert self._maximum_timer_interval is not None
        if current_utc_datetime + self._maximum_timer_interval < self._final_fire_at:
            return current_utc_datetime + self._maximum_timer_interval
        return self._final_fire_at


class WhenAnyTask(CompositeTask[Task[T]], Generic[T]):
    """A task that completes when any of its child tasks complete."""

    def __init__(self, tasks: list[Task[T]]):
        super().__init__(cast(list[Task[Any]], tasks))

    def on_child_completed(self, task: Task[Any]) -> None:
        # The first task to complete is the result of the WhenAnyTask.
        if not self.is_complete:
            self._is_complete = True
            self._result = cast(Task[T], task)
            self._notify_parent()


def when_all(tasks: list[Task[T]]) -> WhenAllTask[T]:
    """Wait for every provided task to complete or fail.

    Results preserve input order and duplicate task references. Tasks may be
    shared between groups without scheduling additional work. If any task fails,
    the first observed failure is raised after every task has finished.
    """
    return WhenAllTask(tasks)


def when_any(tasks: Sequence[Task[T]]) -> WhenAnyTask[T]:
    """Return the first direct child to complete or fail.

    Tasks may be shared between groups. Losing tasks are not canceled.
    """
    return WhenAnyTask(list(tasks))


class ActivityContext:
    def __init__(self, orchestration_id: str, task_id: int):
        self._orchestration_id = orchestration_id
        self._task_id = task_id

    @property
    def orchestration_id(self) -> str:
        """Get the ID of the orchestration instance that scheduled this activity.

        Returns
        -------
        str
            The ID of the current orchestration instance.
        """
        return self._orchestration_id

    @property
    def task_id(self) -> int:
        """Get the task ID associated with this activity invocation.

        The task ID is an auto-incrementing integer that is unique within
        the scope of the orchestration instance. It can be used to distinguish
        between multiple activity invocations that are part of the same
        orchestration instance.

        Returns
        -------
        str
            The ID of the current orchestration instance.
        """
        return self._task_id


# Orchestrators are generators that yield tasks, receive any type, and return TOutput
Orchestrator: TypeAlias = Callable[[OrchestrationContext, TInput], Generator[Task[Any], Any, TOutput] | TOutput]

# Activities are simple functions that can be scheduled by orchestrators
Activity: TypeAlias = Callable[[ActivityContext, TInput], TOutput]

Entity: TypeAlias = Callable[[EntityContext, TInput], TOutput] | type[DurableEntity]


class RetryPolicy:
    """Represents the retry policy for an orchestration or activity function."""

    def __init__(self, *,
                 first_retry_interval: timedelta,
                 max_number_of_attempts: int,
                 backoff_coefficient: float | None = 1.0,
                 max_retry_interval: timedelta | None = None,
                 retry_timeout: timedelta | None = None):
        """Creates a new RetryPolicy instance.

        Parameters
        ----------
        first_retry_interval : timedelta
            The retry interval to use for the first retry attempt.
        max_number_of_attempts : int
            The maximum number of retry attempts.
        backoff_coefficient : float | None
            The backoff coefficient to use for calculating the next retry interval.
        max_retry_interval : timedelta | None
            The maximum retry interval to use for any retry attempt.
        retry_timeout : timedelta | None
            The maximum amount of time to spend retrying the operation.
        """
        # validate inputs
        if first_retry_interval < timedelta(seconds=0):
            raise ValueError('first_retry_interval must be >= 0')
        if max_number_of_attempts < 1:
            raise ValueError('max_number_of_attempts must be >= 1')
        if backoff_coefficient is not None and backoff_coefficient < 1:
            raise ValueError('backoff_coefficient must be >= 1')
        if max_retry_interval is not None and max_retry_interval < timedelta(seconds=0):
            raise ValueError('max_retry_interval must be >= 0')
        if retry_timeout is not None and retry_timeout < timedelta(seconds=0):
            raise ValueError('retry_timeout must be >= 0')

        self._first_retry_interval = first_retry_interval
        self._max_number_of_attempts = max_number_of_attempts
        self._backoff_coefficient = backoff_coefficient
        self._max_retry_interval = max_retry_interval
        self._retry_timeout = retry_timeout

    @property
    def first_retry_interval(self) -> timedelta:
        """The retry interval to use for the first retry attempt."""
        return self._first_retry_interval

    @property
    def max_number_of_attempts(self) -> int:
        """The maximum number of retry attempts."""
        return self._max_number_of_attempts

    @property
    def backoff_coefficient(self) -> float | None:
        """The backoff coefficient to use for calculating the next retry interval."""
        return self._backoff_coefficient

    @property
    def max_retry_interval(self) -> timedelta | None:
        """The maximum retry interval to use for any retry attempt."""
        return self._max_retry_interval

    @property
    def retry_timeout(self) -> timedelta | None:
        """The maximum amount of time to spend retrying the operation."""
        return self._retry_timeout


def get_entity_name(fn: Entity[Any, Any]) -> str:
    if hasattr(fn, "__durable_entity_name__"):
        return getattr(fn, "__durable_entity_name__")
    if isinstance(fn, type) and issubclass(fn, DurableEntity):
        return fn.__name__
    return get_name(cast(Callable[..., Any], fn))


def get_name(fn: Callable[..., Any]) -> str:
    """Returns the name of the provided function"""
    name = fn.__name__
    if name == '<lambda>':
        raise ValueError('Cannot infer a name from a lambda function. Please provide a name explicitly.')

    return name