durabletask 1.3.0.dev27__py3-none-any.whl → 1.3.0.dev29__py3-none-any.whl

durabletask/__init__.py

@@ -4,13 +4,24 @@
 """Durable Task SDK for Python"""
 
 from durabletask.payload.store import LargePayloadStorageOptions, PayloadStore
-from durabletask.worker import ConcurrencyOptions, VersioningOptions
+from durabletask.worker import (
+    ActivityWorkItemFilter,
+    ConcurrencyOptions,
+    EntityWorkItemFilter,
+    OrchestrationWorkItemFilter,
+    VersioningOptions,
+    WorkItemFilters,
+)
 
 __all__ = [
+    "ActivityWorkItemFilter",
     "ConcurrencyOptions",
+    "EntityWorkItemFilter",
     "LargePayloadStorageOptions",
+    "OrchestrationWorkItemFilter",
     "PayloadStore",
     "VersioningOptions",
+    "WorkItemFilters",
 ]
 
 PACKAGE_NAME = "durabletask"

durabletask/task.py

@@ -98,7 +98,7 @@ class OrchestrationContext(ABC):
         pass
 
     @abstractmethod
-    def create_timer(self, fire_at: Union[datetime, timedelta]) -> Task:
+    def create_timer(self, fire_at: Union[datetime, timedelta]) -> CancellableTask:
         """Create a Timer Task to fire after at the specified deadline.
 
         Parameters
@@ -228,10 +228,10 @@ class OrchestrationContext(ABC):
         """
         pass
 
-    # TOOD: Add a timeout parameter, which allows the task to be canceled if the event is
+    # 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.
     @abstractmethod
-    def wait_for_external_event(self, name: str) -> CompletableTask:
+    def wait_for_external_event(self, name: str) -> CancellableTask:
         """Wait asynchronously for an event to be raised with the name `name`.
 
         Parameters
@@ -324,6 +324,10 @@ class OrchestrationStateError(Exception):
     pass
 
 
+class TaskCancelledError(Exception):
+    """Exception type for cancelled orchestration tasks."""
+
+
 class Task(ABC, Generic[T]):
     """Abstract base class for asynchronous tasks in a durable orchestration."""
     _result: T
@@ -435,6 +439,48 @@ class CompletableTask(Task[T]):
             self._parent.on_child_completed(self)
 
 
+class CancellableTask(CompletableTask[T]):
+    """A completable task that can be cancelled before it finishes."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._is_cancelled = False
+        self._cancel_handler: Optional[Callable[[], 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
+        if self._parent is not None:
+            self._parent.on_child_completed(self)
+        return True
+
+
 class RetryableTask(CompletableTask[T]):
     """A task that can be retried according to a retry policy."""
 
@@ -474,14 +520,29 @@ class RetryableTask(CompletableTask[T]):
         return None
 
 
-class TimerTask(CompletableTask[T]):
-
-    def __init__(self) -> None:
+class TimerTask(CancellableTask[None]):
+    def __init__(self, final_fire_at: Optional[datetime] = None,
+                 maximum_timer_interval: Optional[timedelta] = 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):
         self._retryable_parent = retryable_task
 
+    def _handle_timer_fired(self, current_utc_datetime: datetime) -> Optional[datetime]:
+        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:
+        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]):
     """A task that completes when any of its child tasks complete."""

durabletask/testing/in_memory_backend.py

@@ -26,6 +26,7 @@ from google.protobuf import empty_pb2, timestamp_pb2, wrappers_pb2
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.helpers as helpers
+from durabletask.entities.entity_instance_id import EntityInstanceId
 
 
 @dataclass
@@ -56,6 +57,7 @@ class ActivityWorkItem:
     task_id: int
     input: Optional[str]
     completion_token: int
+    version: Optional[str] = None
 
 
 @dataclass
@@ -451,9 +453,57 @@ class InMemoryOrchestrationBackend(stubs.TaskHubSidecarServiceServicer):
             f"Restarted instance '{request.instanceId}' as '{new_instance_id}'")
         return pb.RestartInstanceResponse(instanceId=new_instance_id)
 
+    @staticmethod
+    def _parse_work_item_filters(request: pb.GetWorkItemsRequest):
+        """Extract filters from the request.
+
+        Returns a tuple of three values, one per work-item category.  Each
+        value is either ``None`` (no filtering -- dispatch everything) or a
+        ``dict`` mapping a task name to a ``frozenset`` of accepted versions
+        (empty frozenset means *any* version of that name is accepted).
+        An empty ``dict`` means the worker opted into filtering for that
+        category but listed no names, so *nothing* should match.
+        """
+        if not request.HasField("workItemFilters"):
+            return None, None, None
+        wf = request.workItemFilters
+
+        def _build_filter(filters):
+            result: dict[str, frozenset[str]] = {}
+            for f in filters:
+                versions = frozenset(f.versions) if f.versions else frozenset()
+                existing = result.get(f.name, frozenset())
+                result[f.name] = existing | versions
+            return result
+
+        orch_filter = _build_filter(wf.orchestrations)
+        activity_filter = _build_filter(wf.activities)
+        entity_filter = {f.name: frozenset() for f in wf.entities}
+        return orch_filter, activity_filter, entity_filter
+
+    @staticmethod
+    def _matches_filter(name: str, version: Optional[str],
+                        filt: Optional[dict[str, frozenset[str]]]) -> bool:
+        """Check whether a work item matches the parsed filter.
+
+        *filt* is ``None`` when the worker did not opt into filtering
+        (everything matches).  Otherwise it is a dict mapping accepted
+        names to a frozenset of accepted versions.  An empty frozenset
+        means any version of that name is accepted.
+        """
+        if filt is None:
+            return True
+        accepted_versions = filt.get(name)
+        if accepted_versions is None:
+            return False
+        if not accepted_versions:
+            return True  # empty set -- any version
+        return (version or "") in accepted_versions
+
     def GetWorkItems(self, request: pb.GetWorkItemsRequest, context):
         """Streams work items to the worker (orchestration and activity work items)."""
         self._logger.info("Worker connected and requesting work items")
+        orch_filter, activity_filter, entity_filter = self._parse_work_item_filters(request)
 
         try:
             while context.is_active() and not self._shutdown_event.is_set():
@@ -461,6 +511,7 @@ class InMemoryOrchestrationBackend(stubs.TaskHubSidecarServiceServicer):
 
                 with self._lock:
                     # Check for orchestration work
+                    skipped_orchs: list[str] = []
                     while self._orchestration_queue:
                         instance_id = self._orchestration_queue.popleft()
                         self._orchestration_queue_set.discard(instance_id)
@@ -469,11 +520,15 @@ class InMemoryOrchestrationBackend(stubs.TaskHubSidecarServiceServicer):
                         if not instance or not instance.pending_events:
                             continue
 
+                        # Skip if orchestration doesn't match filters
+                        if not self._matches_filter(
+                                instance.name, instance.version, orch_filter):
+                            skipped_orchs.append(instance_id)
+                            continue
+
                         if instance_id in self._orchestration_in_flight:
                             # Already being processed — re-add to queue
-                            if instance_id not in self._orchestration_queue_set:
-                                self._orchestration_queue.append(instance_id)
-                                self._orchestration_queue_set.add(instance_id)
+                            skipped_orchs.append(instance_id)
                             break
 
                         # Move pending events to dispatched_events
@@ -500,27 +555,66 @@ class InMemoryOrchestrationBackend(stubs.TaskHubSidecarServiceServicer):
                         )
                         break
 
+                    # Re-queue skipped orchestrations for other workers
+                    for s in skipped_orchs:
+                        if s not in self._orchestration_queue_set:
+                            self._orchestration_queue.append(s)
+                            self._orchestration_queue_set.add(s)
+
                     # Check for activity work
                     if not work_item and self._activity_queue:
-                        activity = self._activity_queue.popleft()
-                        work_item = pb.WorkItem(
-                            completionToken=str(activity.completion_token),
-                            activityRequest=pb.ActivityRequest(
-                                name=activity.name,
-                                taskId=activity.task_id,
-                                input=wrappers_pb2.StringValue(value=activity.input) if activity.input else None,
-                                orchestrationInstance=pb.OrchestrationInstance(instanceId=activity.instance_id)
+                        # Scan for the first matching activity
+                        skipped: list = []
+                        matched_activity = None
+                        while self._activity_queue:
+                            candidate = self._activity_queue.popleft()
+                            if not self._matches_filter(
+                                    candidate.name, candidate.version,
+                                    activity_filter):
+                                skipped.append(candidate)
+                                continue
+                            matched_activity = candidate
+                            break
+                        # Put back non-matching items
+                        for s in skipped:
+                            self._activity_queue.append(s)
+
+                        if matched_activity is not None:
+                            work_item = pb.WorkItem(
+                                completionToken=str(matched_activity.completion_token),
+                                activityRequest=pb.ActivityRequest(
+                                    name=matched_activity.name,
+                                    taskId=matched_activity.task_id,
+                                    input=wrappers_pb2.StringValue(value=matched_activity.input) if matched_activity.input else None,
+                                    orchestrationInstance=pb.OrchestrationInstance(instanceId=matched_activity.instance_id)
+                                )
                             )
-                        )
 
                     # Check for entity work
                     if not work_item:
+                        skipped_entities: list[str] = []
                         while self._entity_queue:
                             entity_id = self._entity_queue.popleft()
                             self._entity_queue_set.discard(entity_id)
                             entity = self._entities.get(entity_id)
 
                             if entity and entity.pending_operations:
+                                # Skip if entity name doesn't match filters
+                                if entity_filter is not None:
+                                    try:
+                                        parsed = EntityInstanceId.parse(entity_id)
+                                        if not self._matches_filter(
+                                                parsed.entity, None,
+                                                entity_filter):
+                                            skipped_entities.append(entity_id)
+                                            continue
+                                    except ValueError:
+                                        self._logger.warning(
+                                            f"Cannot parse entity ID '{entity_id}' "
+                                            f"for filter matching; skipping")
+                                        skipped_entities.append(entity_id)
+                                        continue
+
                                 # Skip if this entity is already being processed
                                 if entity_id in self._entity_in_flight:
                                     continue
@@ -547,6 +641,12 @@ class InMemoryOrchestrationBackend(stubs.TaskHubSidecarServiceServicer):
                                 )
                                 break
 
+                        # Re-queue skipped entities for other workers
+                        for s in skipped_entities:
+                            if s not in self._entity_queue_set:
+                                self._entity_queue.append(s)
+                                self._entity_queue_set.add(s)
+
                 if work_item:
                     yield work_item
                 else:
@@ -1274,12 +1374,15 @@ class InMemoryOrchestrationBackend(stubs.TaskHubSidecarServiceServicer):
             instance.status = pb.ORCHESTRATION_STATUS_RUNNING
 
         # Queue activity for execution
+        task_version = schedule_task.version.value \
+            if schedule_task.HasField("version") else None
         self._activity_queue.append(ActivityWorkItem(
             instance_id=instance.instance_id,
             name=task_name,
             task_id=task_id,
             input=input_value,
-            completion_token=instance.completion_token
+            completion_token=instance.completion_token,
+            version=task_version,
         ))
         self._work_available.set()
 

durabletask/worker.py

@@ -9,11 +9,12 @@ import os
 import random
 import time
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 from threading import Event, Thread
 from types import GeneratorType
 from enum import Enum
-from typing import Any, Generator, Optional, Sequence, Tuple, TypeVar, Union
+from typing import Any, Generator, Optional, Sequence, Tuple, TypeVar, Union, overload
 import uuid
 from packaging.version import InvalidVersion, parse
 
@@ -42,6 +43,7 @@ from durabletask.payload.store import PayloadStore
 TInput = TypeVar("TInput")
 TOutput = TypeVar("TOutput")
 DATETIME_STRING_FORMAT = '%Y-%m-%dT%H:%M:%S.%fZ'
+DEFAULT_MAXIMUM_TIMER_INTERVAL = timedelta(days=3)
 
 
 class ConcurrencyOptions:
@@ -143,6 +145,109 @@ class VersioningOptions:
         self.failure_strategy = failure_strategy
 
 
+# Sentinel object used to distinguish "auto-generate filters" from "clear filters (None)".
+_AUTO_GENERATE_FILTERS = object()
+
+
+@dataclass(frozen=True)
+class OrchestrationWorkItemFilter:
+    """Specifies a filter for orchestration work items."""
+
+    name: str
+    """The name of the orchestration to filter."""
+    versions: list[str] = field(default_factory=list)
+    """Optional list of versions to filter."""
+
+
+@dataclass(frozen=True)
+class ActivityWorkItemFilter:
+    """Specifies a filter for activity work items."""
+
+    name: str
+    """The name of the activity to filter."""
+    versions: list[str] = field(default_factory=list)
+    """Optional list of versions to filter."""
+
+
+@dataclass(frozen=True)
+class EntityWorkItemFilter:
+    """Specifies a filter for entity work items.
+
+    The name is normalized to lowercase to match entity registration
+    and instance ID conventions.
+    """
+
+    name: str
+    """The name of the entity to filter."""
+
+    def __post_init__(self):
+        EntityInstanceId.validate_entity_name(self.name)
+        object.__setattr__(self, 'name', self.name.lower())
+
+
+@dataclass(frozen=True)
+class WorkItemFilters:
+    """Work item filters for a Durable Task Worker.
+
+    These filters are passed to the backend and only work items matching the
+    filters will be processed by the worker. If no filters are provided, the
+    worker will process all work items.
+
+    By default, no filters are applied. Call
+    :meth:`TaskHubGrpcWorker.use_work_item_filters` to enable filtering.
+    """
+
+    orchestrations: list[OrchestrationWorkItemFilter] = field(default_factory=list)
+    """List of orchestration filters."""
+    activities: list[ActivityWorkItemFilter] = field(default_factory=list)
+    """List of activity filters."""
+    entities: list[EntityWorkItemFilter] = field(default_factory=list)
+    """List of entity filters."""
+
+    @classmethod
+    def _from_registry(cls, registry: '_Registry') -> 'WorkItemFilters':
+        """Auto-generate work item filters from the task registry."""
+        versions: list[str] = []
+        v = registry.versioning
+        if v and v.match_strategy == VersionMatchStrategy.STRICT and v.version:
+            versions = [registry.versioning.version]
+
+        orchestrations = [
+            OrchestrationWorkItemFilter(name=name, versions=list(versions))
+            for name in registry.orchestrators
+        ]
+        activities = [
+            ActivityWorkItemFilter(name=name, versions=list(versions))
+            for name in registry.activities
+        ]
+        entities = [
+            EntityWorkItemFilter(name=name)
+            for name in registry.entities
+        ]
+        return cls(
+            orchestrations=orchestrations,
+            activities=activities,
+            entities=entities,
+        )
+
+    def _to_grpc(self) -> pb.WorkItemFilters:
+        """Convert to a gRPC WorkItemFilters message."""
+        grpc_filters = pb.WorkItemFilters()
+        for f in self.orchestrations:
+            grpc_filters.orchestrations.append(
+                pb.OrchestrationFilter(name=f.name, versions=f.versions)
+            )
+        for f in self.activities:
+            grpc_filters.activities.append(
+                pb.ActivityFilter(name=f.name, versions=f.versions)
+            )
+        for f in self.entities:
+            grpc_filters.entities.append(
+                pb.EntityFilter(name=f.name)
+            )
+        return grpc_filters
+
+
 class _Registry:
     orchestrators: dict[str, task.Orchestrator]
     activities: dict[str, task.Activity]
@@ -311,7 +416,7 @@ class TaskHubGrpcWorker:
             activity function.
     """
 
-    _response_stream: Optional[grpc.Future] = None
+    _response_stream: Optional[Any] = None
     _interceptors: Optional[list[shared.ClientInterceptor]] = None
 
     def __init__(
@@ -324,6 +429,7 @@ class TaskHubGrpcWorker:
             secure_channel: bool = False,
             interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
             concurrency_options: Optional[ConcurrencyOptions] = None,
+            maximum_timer_interval: Optional[timedelta] = DEFAULT_MAXIMUM_TIMER_INTERVAL,
             payload_store: Optional[PayloadStore] = None,
     ):
         self._registry = _Registry()
@@ -354,12 +460,20 @@ class TaskHubGrpcWorker:
             self._interceptors = None
 
         self._async_worker_manager = _AsyncWorkerManager(self._concurrency_options, self._logger)
+        self._maximum_timer_interval = maximum_timer_interval
+        self._work_item_filters: Optional[WorkItemFilters] = None
+        self._auto_generate_work_item_filters: bool = False
 
     @property
     def concurrency_options(self) -> ConcurrencyOptions:
         """Get the current concurrency options for this worker."""
         return self._concurrency_options
 
+    @property
+    def maximum_timer_interval(self) -> Optional[timedelta]:
+        """Get the configured maximum timer interval for long timer chunking."""
+        return self._maximum_timer_interval
+
     def __enter__(self):
         return self
 
@@ -396,11 +510,65 @@ class TaskHubGrpcWorker:
             raise RuntimeError("Cannot set default version while the worker is running.")
         self._registry.versioning = version
 
+    @overload
+    def use_work_item_filters(self) -> None:
+        ...
+
+    @overload
+    def use_work_item_filters(self, filters: WorkItemFilters) -> None:
+        ...
+
+    @overload
+    def use_work_item_filters(self, filters: None) -> None:
+        ...
+
+    def use_work_item_filters(
+        self,
+        filters: Union[WorkItemFilters, None, object] = _AUTO_GENERATE_FILTERS,
+    ) -> None:
+        """Configures work item filters for the worker.
+
+        Work item filters tell the backend which orchestrations, activities,
+        and entities this worker can handle. When enabled, only matching work
+        items are dispatched to this worker.
+
+        By default no filters are applied and the worker processes all work
+        items. Calling this method enables filtering.
+
+        Args:
+            filters: The filters to apply. If omitted (default), filters are
+                auto-generated from registered orchestrations, activities, and
+                entities at :meth:`start` time. Pass a :class:`WorkItemFilters`
+                instance to provide explicit filters. Pass ``None`` to clear
+                any previously configured filters.
+        """
+        if self._is_running:
+            raise RuntimeError(
+                "Work item filters cannot be changed while the worker is running."
+            )
+        if filters is _AUTO_GENERATE_FILTERS:
+            self._auto_generate_work_item_filters = True
+            self._work_item_filters = None
+        elif filters is None:
+            self._auto_generate_work_item_filters = False
+            self._work_item_filters = None
+        elif isinstance(filters, WorkItemFilters):
+            self._auto_generate_work_item_filters = False
+            self._work_item_filters = filters
+        else:
+            raise TypeError(
+                "filters must be a WorkItemFilters instance, None, or omitted."
+            )
+
     def start(self):
         """Starts the worker on a background thread and begins listening for work items."""
         if self._is_running:
             raise RuntimeError("The worker is already running.")
 
+        # Auto-generate work item filters from registry if opted in
+        if self._auto_generate_work_item_filters:
+            self._work_item_filters = WorkItemFilters._from_registry(self._registry)
+
         def run_loop():
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
@@ -510,6 +678,10 @@ class TaskHubGrpcWorker:
                     maxConcurrentActivityWorkItems=self._concurrency_options.maximum_concurrent_activity_work_items,
                     capabilities=capabilities,
                 )
+                if self._work_item_filters is not None:
+                    get_work_items_request.workItemFilters.CopyFrom(
+                        self._work_item_filters._to_grpc()
+                    )
                 self._response_stream = stub.GetWorkItems(get_work_items_request)
                 self._logger.info(
                     f"Successfully connected to {self._host_address}. Waiting for work items..."
@@ -522,7 +694,11 @@ class TaskHubGrpcWorker:
 
                 def stream_reader():
                     try:
-                        for work_item in self._response_stream:
+                        response_stream = self._response_stream
+                        if response_stream is None:
+                            return
+
+                        for work_item in response_stream:
                             work_item_queue.put(work_item)
                     except Exception as e:
                         work_item_queue.put(e)
@@ -674,7 +850,8 @@ class TaskHubGrpcWorker:
         try:
             executor = _OrchestrationExecutor(
                 self._registry, self._logger,
-                persisted_orch_span_id=persisted_orch_span_id)
+                persisted_orch_span_id=persisted_orch_span_id,
+                maximum_timer_interval=self.maximum_timer_interval)
             result = executor.execute(instance_id, req.pastEvents, req.newEvents)
 
             # Determine completion status for span
@@ -962,7 +1139,11 @@ class _RuntimeOrchestrationContext(task.OrchestrationContext):
     _generator: Optional[Generator[task.Task, Any, Any]]
     _previous_task: Optional[task.Task]
 
-    def __init__(self, instance_id: str, registry: _Registry):
+    def __init__(self,
+                 instance_id: str,
+                 registry: _Registry,
+                 maximum_timer_interval: Optional[timedelta] = DEFAULT_MAXIMUM_TIMER_INTERVAL,
+                 ):
         self._generator = None
         self._is_replaying = True
         self._is_complete = False
@@ -983,12 +1164,13 @@ class _RuntimeOrchestrationContext(task.OrchestrationContext):
         self._version: Optional[str] = None
         self._completion_status: Optional[pb.OrchestrationStatus] = None
         self._received_events: dict[str, list[Any]] = {}
-        self._pending_events: dict[str, list[task.CompletableTask]] = {}
+        self._pending_events: dict[str, list[task.CancellableTask]] = {}
         self._new_input: Optional[Any] = None
         self._save_events = False
         self._encoded_custom_status: Optional[str] = None
         self._parent_trace_context: Optional[pb.TraceContext] = None
         self._orchestration_trace_context: Optional[pb.TraceContext] = None
+        self._maximum_timer_interval = maximum_timer_interval
 
     def run(self, generator: Generator[task.Task, Any, Any]):
         self._generator = generator
@@ -1154,7 +1336,7 @@ class _RuntimeOrchestrationContext(task.OrchestrationContext):
             shared.to_json(custom_status) if custom_status is not None else None
         )
 
-    def create_timer(self, fire_at: Union[datetime, timedelta]) -> task.Task:
+    def create_timer(self, fire_at: Union[datetime, timedelta]) -> task.CancellableTask:
         return self.create_timer_internal(fire_at)
 
     def create_timer_internal(
@@ -1164,11 +1346,30 @@ class _RuntimeOrchestrationContext(task.OrchestrationContext):
     ) -> task.TimerTask:
         id = self.next_sequence_number()
         if isinstance(fire_at, timedelta):
-            fire_at = self.current_utc_datetime + fire_at
-        action = ph.new_create_timer_action(id, fire_at)
+            final_fire_at = self.current_utc_datetime + fire_at
+        else:
+            final_fire_at = fire_at
+
+        next_fire_at: datetime = final_fire_at
+
+        if (
+            self._maximum_timer_interval is not None
+            and self._maximum_timer_interval > timedelta(0)
+            and self.current_utc_datetime + self._maximum_timer_interval < final_fire_at
+        ):
+            timer_task = task.TimerTask(final_fire_at, self._maximum_timer_interval)
+            next_fire_at = timer_task._get_next_fire_at(self.current_utc_datetime)
+        else:
+            timer_task = task.TimerTask()
+
+        action = ph.new_create_timer_action(id, next_fire_at)
         self._pending_actions[id] = action
 
-        timer_task: task.TimerTask = task.TimerTask()
+        def _cancel_timer() -> None:
+            self._pending_actions.pop(id, None)
+            self._pending_tasks.pop(id, None)
+
+        timer_task.set_cancel_handler(_cancel_timer)
         if retryable_task is not None:
             timer_task.set_retryable_parent(retryable_task)
         self._pending_tasks[id] = timer_task
@@ -1399,13 +1600,13 @@ class _RuntimeOrchestrationContext(task.OrchestrationContext):
             action = pb.OrchestratorAction(id=task_id, sendEntityMessage=entity_unlock_message)
             self._pending_actions[task_id] = action
 
-    def wait_for_external_event(self, name: str) -> task.CompletableTask:
+    def wait_for_external_event(self, name: str) -> task.CancellableTask:
         # Check to see if this event has already been received, in which case we
         # can return it immediately. Otherwise, record out intent to receive an
         # event with the given name so that we can resume the generator when it
         # arrives. If there are multiple events with the same name, we return
         # them in the order they were received.
-        external_event_task: task.CompletableTask = task.CompletableTask()
+        external_event_task: task.CancellableTask = task.CancellableTask()
         event_name = name.casefold()
         event_list = self._received_events.get(event_name, None)
         if event_list:
@@ -1419,6 +1620,19 @@ class _RuntimeOrchestrationContext(task.OrchestrationContext):
                 task_list = []
                 self._pending_events[event_name] = task_list
             task_list.append(external_event_task)
+
+            def _cancel_wait() -> None:
+                waiting_tasks = self._pending_events.get(event_name)
+                if waiting_tasks is None:
+                    return
+                try:
+                    waiting_tasks.remove(external_event_task)
+                except ValueError:
+                    return
+                if not waiting_tasks:
+                    del self._pending_events[event_name]
+
+            external_event_task.set_cancel_handler(_cancel_wait)
         return external_event_task
 
     def continue_as_new(self, new_input, *, save_events: bool = False) -> None:
@@ -1461,9 +1675,11 @@ class _OrchestrationExecutor:
         registry: _Registry,
         logger: logging.Logger,
         persisted_orch_span_id: Optional[str] = None,
+        maximum_timer_interval: Optional[timedelta] = DEFAULT_MAXIMUM_TIMER_INTERVAL,
     ):
         self._registry = registry
         self._logger = logger
+        self._maximum_timer_interval = maximum_timer_interval
         self._is_suspended = False
         self._suspended_events: list[pb.HistoryEvent] = []
         self._persisted_orch_span_id = persisted_orch_span_id
@@ -1497,7 +1713,11 @@ class _OrchestrationExecutor:
                 "The new history event list must have at least one event in it."
             )
 
-        ctx = _RuntimeOrchestrationContext(instance_id, self._registry)
+        ctx = _RuntimeOrchestrationContext(
+            instance_id,
+            self._registry,
+            maximum_timer_interval=self._maximum_timer_interval,
+        )
         try:
             # Rebuild local state by replaying old history into the orchestrator function
             self._logger.debug(
@@ -1654,6 +1874,12 @@ class _OrchestrationExecutor:
                             f"{ctx.instance_id}: Ignoring unexpected timerFired event with ID = {timer_id}."
                         )
                     return
+                if not isinstance(timer_task, task.TimerTask):
+                    if not ctx._is_replaying:
+                        self._logger.warning(
+                            f"{ctx.instance_id}: Ignoring timerFired event with non-timer task ID = {timer_id}."
+                        )
+                    return
                 # Emit timer span with backdated start time (skip during replay)
                 if not ctx.is_replaying:
                     timer_info = self._timer_fire_at.get(timer_id)
@@ -1665,27 +1891,39 @@ class _OrchestrationExecutor:
                             scheduled_time_ns=created_ns,
                             parent_trace_context=ctx._orchestration_trace_context or ctx._parent_trace_context,
                         )
-                timer_task.complete(None)
-                if timer_task._retryable_parent is not None:
-                    activity_action = timer_task._retryable_parent._action
+                next_fire_at = timer_task._handle_timer_fired(event.timerFired.fireAt.ToDatetime())
+                if next_fire_at is not None:
+                    id = ctx.next_sequence_number()
+                    new_action = ph.new_create_timer_action(id, next_fire_at)
+                    ctx._pending_tasks[id] = timer_task
+                    ctx._pending_actions[id] = new_action
+
+                    def _cancel_timer() -> None:
+                        ctx._pending_actions.pop(id, None)
+                        ctx._pending_tasks.pop(id, None)
+
+                    timer_task.set_cancel_handler(_cancel_timer)
+                else:
+                    if timer_task._retryable_parent is not None:
+                        activity_action = timer_task._retryable_parent._action
 
-                    if not timer_task._retryable_parent._is_sub_orch:
-                        cur_task = activity_action.scheduleTask
-                        instance_id = None
+                        if not timer_task._retryable_parent._is_sub_orch:
+                            cur_task = activity_action.scheduleTask
+                            instance_id = None
+                        else:
+                            cur_task = activity_action.createSubOrchestration
+                            instance_id = cur_task.instanceId
+                        ctx.call_activity_function_helper(
+                            id=activity_action.id,
+                            activity_function=cur_task.name,
+                            input=cur_task.input.value,
+                            retry_policy=timer_task._retryable_parent._retry_policy,
+                            is_sub_orch=timer_task._retryable_parent._is_sub_orch,
+                            instance_id=instance_id,
+                            fn_task=timer_task._retryable_parent,
+                        )
                     else:
-                        cur_task = activity_action.createSubOrchestration
-                        instance_id = cur_task.instanceId
-                    ctx.call_activity_function_helper(
-                        id=activity_action.id,
-                        activity_function=cur_task.name,
-                        input=cur_task.input.value,
-                        retry_policy=timer_task._retryable_parent._retry_policy,
-                        is_sub_orch=timer_task._retryable_parent._is_sub_orch,
-                        instance_id=instance_id,
-                        fn_task=timer_task._retryable_parent,
-                    )
-                else:
-                    ctx.resume()
+                        ctx.resume()
             elif event.HasField("taskScheduled"):
                 # This history event confirms that the activity execution was successfully scheduled.
                 # Remove the taskScheduled event from the pending action list so we don't schedule it again.

{durabletask-1.3.0.dev27.dist-info → durabletask-1.3.0.dev29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask
-Version: 1.3.0.dev27
+Version: 1.3.0.dev29
 Summary: A Durable Task Client SDK for Python
 License:     MIT License
         

{durabletask-1.3.0.dev27.dist-info → durabletask-1.3.0.dev29.dist-info}/RECORD

@@ -1,8 +1,8 @@
-durabletask/__init__.py,sha256=Xb-2zUIwDbCCl4Q_89TqMJ-3zckHFrGtSMcVhEeSseQ,407
+durabletask/__init__.py,sha256=OdfKCNlS_NJawRfLWsFNj7YIHeGSQkh2VH3OzG0Oric,644
 durabletask/client.py,sha256=NbIdDTQR7XI_ZiqsGMP0q5vmbe5-ShyGUQre1qgB-Ag,34107
 durabletask/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-durabletask/task.py,sha256=YeZj-tYHiLml672hAMtri8_hqnpmR5in4AHMi5owFQo,21793
-durabletask/worker.py,sha256=TNc1WJTmh42sLWyD_FuwUw7YqrIPA7AFpJKd0iHOn5M,124479
+durabletask/task.py,sha256=QS7RPUaWhb4Xq5hi7nnueooJqQSDO6ipEoiAiSVkAIs,24009
+durabletask/worker.py,sha256=B8aYoOguNuYccmmtGF9GFr3YEEdfTGJqcOoox0kByQM,133661
 durabletask/entities/__init__.py,sha256=DbNd5riqWZaj3tG6gN82O8Q6wTmFpe6QaH0pQgDSPHs,721
 durabletask/entities/durable_entity.py,sha256=LQPWnUlRsHiFVRoTdpeSK--eXtjf2UGbVQwEEKf7QwI,3318
 durabletask/entities/entity_context.py,sha256=U-B3i9QP34N-6Fikx_tMp8zo0YLdmwhwpdxwjHd7z-M,5346
@@ -31,9 +31,9 @@ durabletask/payload/__init__.py,sha256=1h68pQvgk8JUp5LBJuBq9W4GUPYkdlhqmCCQEg6YB
 durabletask/payload/helpers.py,sha256=RYG5MEVAqHjm4zfFHs3Td91FVQHUoCcb5hbEJ4sYj5s,12350
 durabletask/payload/store.py,sha256=3qJMvKxRUkr6ScWUzxpKAVgzuhFLywRW8a2_5OOmNk4,3000
 durabletask/testing/__init__.py,sha256=rXbcSFtzuaRAbDNX-HmdgbxLTegvKJ1FRjZfSOIAMgA,323
-durabletask/testing/in_memory_backend.py,sha256=REmzhgAAw_AOpxrRAEbZlU4jqQLZo6QdbWYoWaruBDo,74102
-durabletask-1.3.0.dev27.dist-info/licenses/LICENSE,sha256=ws_MuBL-SCEBqPBFl9_FqZkaaydIJmxHrJG2parhU4M,1141
-durabletask-1.3.0.dev27.dist-info/METADATA,sha256=uGCKNOJQtqqay5ULKwLclHtAQh3zHYj39_ZP37C1lMU,4404
-durabletask-1.3.0.dev27.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-durabletask-1.3.0.dev27.dist-info/top_level.txt,sha256=EBVyuKWnjOwq8bJI1Uvb9U3c4fzQxACWj9p83he6fik,12
-durabletask-1.3.0.dev27.dist-info/RECORD,,
+durabletask/testing/in_memory_backend.py,sha256=ELxyCDRDNOabygIvw9ZeRUpP3MzeM5Hdbu6QlRwKdno,79249
+durabletask-1.3.0.dev29.dist-info/licenses/LICENSE,sha256=ws_MuBL-SCEBqPBFl9_FqZkaaydIJmxHrJG2parhU4M,1141
+durabletask-1.3.0.dev29.dist-info/METADATA,sha256=NFWHPGrcuEDhGp2-ut4FlmAyjtBdlOhTMz9ZuCb_EOI,4404
+durabletask-1.3.0.dev29.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+durabletask-1.3.0.dev29.dist-info/top_level.txt,sha256=EBVyuKWnjOwq8bJI1Uvb9U3c4fzQxACWj9p83he6fik,12
+durabletask-1.3.0.dev29.dist-info/RECORD,,