Add rewind support (#121)

* Initial implementation

* Fix in-memory-backend history ordering, rewind check

* Match other rewind implementations better

* Lint + pylance

* Lint

* PR Feedback

* Add DTS tests

* Skip DTS E2E rewind tests until available

* Add warning

* Add more tests, fix in-memory backend bug

* Restore PROTO_SOURCE_COMMIT_HASH

* Remove duplicate test

* Add RewindInstance to _SyncTaskHubSidecarServiceStub protocol for pyright

* Address PR review feedback; enable DTS rewind e2e tests

- Add async rewind_orchestration parity on AsyncTaskHubGrpcClient

- in-memory backend: stamp ExecutionRewound with current time; avoid duplicate executionRewound on post-rewind replay

- worker: fix rewind comments, guard parentInstance, drop dead sub-orchestration branch

- tests: free-port fixture, clarify docstrings, un-skip DTS rewind e2e (emulator now supports RewindInstance)

- azuremanaged CHANGELOG entry

* Clarify rewind jump-start comment (PR review nit)

Author: andystaples

CHANGELOG.md

@@ -7,6 +7,10 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
+ADDED
+
+- Added `TaskHubGrpcClient.rewind_orchestration()` to rewind a failed orchestration instance to its last known good state. Failed activity and sub-orchestration results are removed from the history and the orchestration replays from the last successful checkpoint, retrying only the failed work. The in-memory testing backend supports rewind as well.
+
 ## v1.7.2
 
 FIXED

durabletask-azuremanaged/CHANGELOG.md

@@ -7,6 +7,8 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
+- Added `rewind_orchestration()` to `DurableTaskSchedulerClient` and `AsyncDurableTaskSchedulerClient` (inherited from the base clients) to rewind a failed orchestration instance to its last known good state.
+
 ## v1.7.2
 
 - Updates base dependency to durabletask v1.7.2.

durabletask/client.py

@@ -320,6 +320,9 @@ def SuspendInstance(self, request: pb.SuspendRequest) -> pb.SuspendResponse:
     def ResumeInstance(self, request: pb.ResumeRequest) -> pb.ResumeResponse:
         ...
 
+    def RewindInstance(self, request: pb.RewindInstanceRequest) -> pb.RewindInstanceResponse:
+        ...
+
     def RestartInstance(self, request: pb.RestartInstanceRequest) -> pb.RestartInstanceResponse:
         ...
 
@@ -381,6 +384,9 @@ async def SuspendInstance(self, request: pb.SuspendRequest) -> pb.SuspendRespons
     async def ResumeInstance(self, request: pb.ResumeRequest) -> pb.ResumeResponse:
         ...
 
+    async def RewindInstance(self, request: pb.RewindInstanceRequest) -> pb.RewindInstanceResponse:
+        ...
+
     async def RestartInstance(self, request: pb.RestartInstanceRequest) -> pb.RestartInstanceResponse:
         ...
 
@@ -782,6 +788,26 @@ def resume_orchestration(self, instance_id: str) -> None:
         self._logger.info(f"Resuming instance '{instance_id}'.")
         self._stub.ResumeInstance(req)
 
+    def rewind_orchestration(self, instance_id: str, *,
+                             reason: str | None = None) -> None:
+        """Rewinds a failed orchestration instance to its last known good state.
+
+        Rewind removes failed task and sub-orchestration results from the
+        orchestration history and replays the orchestration from the last
+        successful checkpoint.  Activities that previously succeeded are
+        not re-executed; only failed work is retried.
+
+        Args:
+            instance_id: The ID of the orchestration instance to rewind.
+            reason: An optional reason string describing why the orchestration is being rewound.
+        """
+        req = pb.RewindInstanceRequest(
+            instanceId=instance_id,
+            reason=helpers.get_string_value(reason))
+
+        self._logger.info(f"Rewinding instance '{instance_id}'.")
+        self._stub.RewindInstance(req)
+
     def restart_orchestration(self, instance_id: str, *,
                               restart_with_new_instance_id: bool = False) -> str:
         """Restarts an existing orchestration instance.
@@ -1267,6 +1293,26 @@ async def resume_orchestration(self, instance_id: str) -> None:
         self._logger.info(f"Resuming instance '{instance_id}'.")
         await self._stub.ResumeInstance(req)
 
+    async def rewind_orchestration(self, instance_id: str, *,
+                                   reason: str | None = None) -> None:
+        """Rewinds a failed orchestration instance to its last known good state.
+
+        Rewind removes failed task and sub-orchestration results from the
+        orchestration history and replays the orchestration from the last
+        successful checkpoint.  Activities that previously succeeded are
+        not re-executed; only failed work is retried.
+
+        Args:
+            instance_id: The ID of the orchestration instance to rewind.
+            reason: An optional reason string describing why the orchestration is being rewound.
+        """
+        req = pb.RewindInstanceRequest(
+            instanceId=instance_id,
+            reason=helpers.get_string_value(reason))
+
+        self._logger.info(f"Rewinding instance '{instance_id}'.")
+        await self._stub.RewindInstance(req)
+
     async def restart_orchestration(self, instance_id: str, *,
                                     restart_with_new_instance_id: bool = False) -> str:
         """Restarts an existing orchestration instance.

durabletask/internal/helpers.py

@@ -236,6 +236,21 @@ def new_terminated_event(*, encoded_output: str | None = None) -> pb.HistoryEven
     )
 
 
+def new_execution_completed_event(
+        status: pb.OrchestrationStatus,
+        encoded_result: str | None = None,
+        failure_details: pb.TaskFailureDetails | None = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionCompleted=pb.ExecutionCompletedEvent(
+            orchestrationStatus=status,
+            result=get_string_value(encoded_result),
+            failureDetails=failure_details,
+        )
+    )
+
+
 def get_string_value(val: str | None) -> wrappers_pb2.StringValue | None:
     if val is None:
         return None

durabletask/testing/README.md

@@ -255,8 +255,7 @@ The in-memory backend is designed for testing and has some limitations compared
 1. **No persistence**: All state is lost when the backend is stopped
 2. **No distributed execution**: Runs in a single process
 3. **No history streaming**: StreamInstanceHistory is not implemented
-4. **No rewind**: RewindInstance is not implemented
-5. **No recursive termination**: Recursive termination is not supported
+4. **No recursive termination**: Recursive termination is not supported
 
 ### Best Practices
 

durabletask/testing/in_memory_backend.py

@@ -273,14 +273,16 @@ def StartInstance(self, request: pb.CreateInstanceRequest, context: grpc.Service
             )
             self._next_completion_token += 1
 
-            # Add initial events to start the orchestration
-            orchestrator_started = helpers.new_orchestrator_started_event(start_time)
+            # Add initial events to start the orchestration.
+            # orchestratorStarted bookends each replay batch and is
+            # always the very first event, followed by executionStarted.
             execution_started = helpers.new_execution_started_event(
                 request.name, instance_id,
                 request.input.value if request.input else None,
                 dict(request.tags) if request.tags else None,
                 version=version,
             )
+            orchestrator_started = helpers.new_orchestrator_started_event(start_time)
 
             instance.pending_events.append(orchestrator_started)
             instance.pending_events.append(execution_started)
@@ -611,8 +613,12 @@ def GetWorkItems(self, request: pb.GetWorkItemsRequest, context: grpc.ServicerCo
                         instance.pending_events.clear()
 
                         # Add OrchestratorStarted for re-dispatches so that
-                        # ctx.current_utc_datetime advances correctly
-                        if instance.history:
+                        # ctx.current_utc_datetime advances correctly, unless
+                        # the pending batch already begins with one (e.g. a
+                        # post-rewind replay supplies its own).
+                        if instance.history and not (
+                                instance.dispatched_events
+                                and instance.dispatched_events[0].HasField("orchestratorStarted")):
                             now = datetime.now(timezone.utc)
                             orch_started = helpers.new_orchestrator_started_event(now)
                             instance.dispatched_events.insert(0, orch_started)
@@ -802,6 +808,22 @@ def CompleteOrchestratorTask(self, request: pb.OrchestratorResponse, context: gr
             instance.completion_token = self._next_completion_token
             self._next_completion_token += 1
 
+            # Bookend the replay with orchestratorCompleted.
+            # Skip for continue-as-new (status is PENDING after reset).
+            if instance.status != pb.ORCHESTRATION_STATUS_PENDING:
+                instance.history.append(helpers.new_orchestrator_completed_event())
+
+                # executionCompleted is the very last event when the
+                # orchestration reaches a terminal state.
+                if self._is_terminal_status(instance.status):
+                    instance.history.append(
+                        helpers.new_execution_completed_event(
+                            instance.status,
+                            instance.output,
+                            instance.failure_details,
+                        )
+                    )
+
             # Remove from in-flight before notifying or re-enqueuing
             self._orchestration_in_flight.discard(request.instanceId)
 
@@ -1193,8 +1215,33 @@ def DeleteTaskHub(self, request: pb.DeleteTaskHubRequest, context: grpc.Servicer
         return pb.DeleteTaskHubResponse()
 
     def RewindInstance(self, request: pb.RewindInstanceRequest, context: grpc.ServicerContext) -> pb.RewindInstanceResponse:
-        """Rewinds an orchestration instance (not implemented)."""
-        context.abort(grpc.StatusCode.UNIMPLEMENTED, "RewindInstance not implemented")
+        """Rewinds a failed orchestration instance.
+
+        The backend validates the instance is in a failed state, appends
+        an ``ExecutionRewoundEvent`` to the pending events, resets the
+        instance status to RUNNING, and re-enqueues the orchestration
+        so the worker can replay it and produce a
+        ``RewindOrchestrationAction`` with the corrected history.
+        """
+        with self._lock:
+            instance = self._instances.get(request.instanceId)
+            if not instance:
+                context.abort(
+                    grpc.StatusCode.NOT_FOUND,
+                    f"Orchestration instance '{request.instanceId}' not found")
+                return pb.RewindInstanceResponse()
+
+            if instance.status != pb.ORCHESTRATION_STATUS_FAILED:
+                context.abort(
+                    grpc.StatusCode.FAILED_PRECONDITION,
+                    f"Orchestration instance '{request.instanceId}' is not in a failed state")
+                return pb.RewindInstanceResponse()
+
+            reason = request.reason.value if request.HasField("reason") else None
+            self._prepare_rewind(instance, reason)
+
+        self._logger.info(f"Rewound instance '{request.instanceId}'")
+        return pb.RewindInstanceResponse()
 
     def AbandonTaskActivityWorkItem(self, request: pb.AbandonActivityTaskRequest, context: grpc.ServicerContext) -> pb.AbandonActivityTaskResponse:
         """Abandons an activity work item."""
@@ -1280,9 +1327,9 @@ def _create_instance_internal(self, instance_id: str, name: str,
         )
         self._next_completion_token += 1
 
-        orchestrator_started = helpers.new_orchestrator_started_event(now)
         execution_started = helpers.new_execution_started_event(
             name, instance_id, encoded_input, version=version)
+        orchestrator_started = helpers.new_orchestrator_started_event(now)
         instance.pending_events.append(orchestrator_started)
         instance.pending_events.append(execution_started)
 
@@ -1410,14 +1457,25 @@ def _process_action(self, instance: OrchestrationInstance, action: pb.Orchestrat
             self._process_send_event_action(action.sendEvent)
         elif action.HasField("sendEntityMessage"):
             self._process_send_entity_message_action(instance, action)
+        elif action.HasField("rewindOrchestration"):
+            self._process_rewind_orchestration_action(instance, action.rewindOrchestration)
 
     def _process_complete_orchestration_action(self, instance: OrchestrationInstance,
                                                complete_action: pb.CompleteOrchestrationAction):
         """Processes a complete orchestration action."""
         status = complete_action.orchestrationStatus
         instance.status = status
         instance.output = complete_action.result.value if complete_action.result else None
-        instance.failure_details = complete_action.failureDetails if complete_action.failureDetails else None
+        # NOTE: a protobuf singular message field is always truthy, so
+        # ``if complete_action.failureDetails`` would set failure_details
+        # even for successful completions. Use HasField for a correct
+        # presence check so failure_details is None unless the
+        # orchestration actually failed.
+        instance.failure_details = (
+            complete_action.failureDetails
+            if complete_action.HasField("failureDetails")
+            else None
+        )
         instance.completed_at = datetime.now(timezone.utc) if self._is_terminal_status(status) else None
 
         if status == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW:
@@ -1446,11 +1504,11 @@ def _process_complete_orchestration_action(self, instance: OrchestrationInstance
             # Build the new pending events in the correct order:
             # OrchestratorStarted, ExecutionStarted, carryover, new arrivals
             now = datetime.now(timezone.utc)
-            orchestrator_started = helpers.new_orchestrator_started_event(now)
             execution_started = helpers.new_execution_started_event(
                 instance.name, instance.instance_id, new_input,
                 version=new_version,
             )
+            orchestrator_started = helpers.new_orchestrator_started_event(now)
             instance.pending_events.append(orchestrator_started)
             instance.pending_events.append(execution_started)
             instance.pending_events.extend(carryover_events)
@@ -1822,6 +1880,129 @@ def fire() -> None:
         timer_thread = threading.Thread(target=fire, daemon=True)
         timer_thread.start()
 
+    def _prepare_rewind(self, instance: OrchestrationInstance,
+                        reason: str | None = None):
+        """Prepares an orchestration instance for rewind.
+
+        Appends an ``ExecutionRewoundEvent`` to the pending events, resets
+        the instance status to RUNNING, and re-enqueues it so the worker
+        can replay it. The actual history rewriting is done by the SDK
+        worker when it processes the rewind event.
+
+        Args:
+            instance: The orchestration instance to rewind.
+            reason: Optional reason string for the rewind.
+
+        Note:
+            Must be called while holding ``self._lock``.
+        """
+        # Reset instance state so it can be re-processed.
+        instance.status = pb.ORCHESTRATION_STATUS_RUNNING
+        instance.output = None
+        instance.failure_details = None
+        instance.last_updated_at = datetime.now(timezone.utc)
+
+        # Clear any stale dispatched events.
+        instance.dispatched_events.clear()
+
+        # Add the ExecutionRewound event as a new pending event. Stamp it
+        # with the current time (rather than the default epoch) so the
+        # audit event carries a meaningful timestamp, consistent with
+        # other backend-generated events.
+        rewind_timestamp = timestamp_pb2.Timestamp()
+        rewind_timestamp.FromDatetime(datetime.now(timezone.utc))
+        rewind_event = pb.HistoryEvent(
+            eventId=-1,
+            timestamp=rewind_timestamp,
+            executionRewound=pb.ExecutionRewoundEvent(
+                reason=wrappers_pb2.StringValue(value=reason) if reason else None,
+            ),
+        )
+        instance.pending_events.append(rewind_event)
+
+        # Refresh the completion token and enqueue.
+        instance.completion_token = self._next_completion_token
+        self._next_completion_token += 1
+        self._orchestration_in_flight.discard(instance.instance_id)
+        self._enqueue_orchestration(instance.instance_id)
+
+    def _process_rewind_orchestration_action(
+            self, instance: OrchestrationInstance,
+            rewind_action: pb.RewindOrchestrationAction):
+        """Processes a RewindOrchestrationAction returned by the SDK.
+
+        The action contains a ``newHistory`` field with the rewritten
+        history computed by the SDK (failed tasks and sub-orchestration
+        failures removed).  The backend replaces the instance's history
+        with this new history, recursively rewinds any failed
+        sub-orchestrations, and re-enqueues the orchestration.
+        """
+        new_history = list(rewind_action.newHistory)
+
+        # Replace history with the rewritten version.
+        instance.history = new_history
+        instance.status = pb.ORCHESTRATION_STATUS_RUNNING
+        instance.output = None
+        instance.failure_details = None
+        instance.last_updated_at = datetime.now(timezone.utc)
+
+        # Identify sub-orchestrations that were created but did not
+        # complete successfully — they need to be recursively rewound.
+        completed_sub_orch_task_ids: set[int] = set()
+        created_sub_orch_events: dict[int, pb.HistoryEvent] = {}
+        for event in new_history:
+            if event.HasField("subOrchestrationInstanceCreated"):
+                created_sub_orch_events[event.eventId] = event
+            elif event.HasField("subOrchestrationInstanceCompleted"):
+                completed_sub_orch_task_ids.add(
+                    event.subOrchestrationInstanceCompleted.taskScheduledId)
+
+        # Extract the rewind reason from the last ExecutionRewound event.
+        reason: str | None = None
+        for event in reversed(new_history):
+            if event.HasField("executionRewound"):
+                if event.executionRewound.HasField("reason"):
+                    reason = event.executionRewound.reason.value
+                break
+
+        # Recursively rewind failed sub-orchestrations.  If the sub was
+        # purged (no longer in _instances), re-create it from the
+        # subOrchestrationInstanceCreated event so it runs fresh.
+        for task_id, event in created_sub_orch_events.items():
+            if task_id not in completed_sub_orch_task_ids:
+                sub_info = event.subOrchestrationInstanceCreated
+                sub_instance_id = sub_info.instanceId
+                sub_instance = self._instances.get(sub_instance_id)
+                if sub_instance is None:
+                    # Sub-orchestration was purged — re-create it.
+                    sub_name = sub_info.name
+                    sub_input = sub_info.input.value if sub_info.HasField("input") else None
+                    sub_version = sub_info.version.value if sub_info.HasField("version") else None
+                    self._create_instance_internal(
+                        sub_instance_id, sub_name, sub_input, version=sub_version)
+                elif sub_instance.status == pb.ORCHESTRATION_STATUS_FAILED:
+                    self._prepare_rewind(sub_instance, reason)
+                self._watch_sub_orchestration(
+                    instance.instance_id, sub_instance_id, task_id)
+
+        # Re-enqueue so the orchestration replays with the clean history.
+        # The executionRewound event is already present in the clean
+        # history (it was kept by the SDK's rewind rewrite), so it must
+        # NOT be re-sent as a pending event — doing so would duplicate it
+        # in the instance history once the dispatched events are
+        # committed. A lone orchestratorStarted is enough to make the
+        # instance dispatchable for the post-rewind replay; the worker
+        # replays normally because executionCompleted is no longer in the
+        # history.
+        instance.pending_events.clear()
+        instance.dispatched_events.clear()
+        instance.pending_events.append(
+            helpers.new_orchestrator_started_event(datetime.now(timezone.utc)))
+        instance.completion_token = self._next_completion_token
+        self._next_completion_token += 1
+        self._orchestration_in_flight.discard(instance.instance_id)
+        self._enqueue_orchestration(instance.instance_id)
+
     def _enqueue_entity(self, entity_id: str):
         """Enqueues an entity for processing."""
         if entity_id not in self._entity_queue_set: