From 7d5f31dff5947ba7a04e794d670854a143aba8b9 Mon Sep 17 00:00:00 2001
From: Andy Staples <andystaples@microsoft.com>
Date: Wed, 11 Nov 2026 12:34:42 -0700
Subject: [PATCH 1/3] Add coroutine orchestrator dispatch

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 durabletask/internal/asyncio_loop.py          | 225 +++++++++
 durabletask/internal/shared.py                |  18 +
 durabletask/task.py                           |  41 +-
 durabletask/worker.py                         | 128 +++--
 .../test_coroutine_orchestrators.py           | 445 ++++++++++++++++++
 .../test_orchestration_generator_cleanup.py   |   2 +-
 6 files changed, 808 insertions(+), 51 deletions(-)
 create mode 100644 durabletask/internal/asyncio_loop.py
 create mode 100644 tests/durabletask/test_coroutine_orchestrators.py

diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py
new file mode 100644
index 00000000..8865123b
--- /dev/null
+++ b/durabletask/internal/asyncio_loop.py
@@ -0,0 +1,225 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Dispatch-local asyncio scheduling on top of the shared durable history engine."""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import logging
+from collections import deque
+from collections.abc import Callable, Coroutine, Generator
+from datetime import datetime
+from typing import Any, TypeVar
+
+from durabletask import task
+
+T = TypeVar("T")
+
+
+class _CompletionObserver(task.CompositeTask[None]):
+    def __init__(self, durable: task.Task[Any], callback: Callable[[], None]) -> None:
+        self._callback = callback
+        super().__init__([durable])
+
+    def on_child_completed(self, task: task.Task[Any]) -> None:
+        if not self._is_complete:
+            self._is_complete = True
+            self._notify_parent()
+            self._callback()
+
+    def detach(self) -> None:
+        self._notify_parent()
+
+
+class DurableEventLoop(asyncio.AbstractEventLoop):
+    """A FIFO microtask loop; it never polls a selector or advances wall time.
+
+    Only history events can complete durable waits. asyncio owns coroutine
+    stepping; this loop owns the ordering of ready callbacks.
+    """
+
+    def __init__(
+        self, ctx: task.OrchestrationContext, *, is_complete: Callable[[], bool],
+        logger: logging.Logger | None = None,
+    ) -> None:
+        self.ctx = ctx
+        self._is_complete = is_complete
+        self._logger = logger or logging.getLogger(__name__)
+        self._ready: deque[asyncio.Handle] = deque()
+        self._tasks: list[asyncio.Task[Any]] = []
+        self._waits: dict[int, tuple[task.Task[Any], list[asyncio.Future[Any]]]] = {}
+        self._observers: dict[int, _CompletionObserver] = {}
+        self._completions: deque[Callable[[], None]] = deque()
+        self._closed = False
+        self._closing = False
+        self._fatal: BaseException | None = None
+        self._root: asyncio.Task[Any] | None = None
+
+    def time(self) -> float:
+        # A fixed UTC epoch avoids local-time and platform timestamp differences.
+        return (self.ctx.current_utc_datetime - datetime(1970, 1, 1)).total_seconds()
+
+    def get_debug(self) -> bool:
+        return False
+
+    def is_running(self) -> bool:
+        return asyncio._get_running_loop() is self
+
+    def is_closed(self) -> bool:
+        return self._closed
+
+    def create_future(self) -> asyncio.Future[Any]:
+        return asyncio.Future(loop=self)
+
+    def get_task_factory(self) -> None:
+        return None
+
+    def create_task(
+        self, coro: Coroutine[Any, Any, T] | Generator[Any, None, T], *, name: str | None = None,
+        context: contextvars.Context | None = None, **kwargs: Any,
+    ) -> asyncio.Task[T]:
+        if self._closed:
+            coro.close()
+            raise RuntimeError("The durable orchestration dispatch is closed.")
+        if kwargs:
+            coro.close()
+            raise TypeError("Custom/eager task execution is not supported in durable orchestrators.")
+        future = asyncio.Task(coro, loop=self, name=name, context=context)
+        self._tasks.append(future)
+        return future
+
+    def call_soon(
+        self, callback: Callable[..., Any], *args: Any,
+        context: contextvars.Context | None = None,
+    ) -> asyncio.Handle:
+        handle = asyncio.Handle(callback, args, self, context=context)
+        self._ready.append(handle)
+        return handle
+
+    async def wait_durable(self, durable: task.Task[T]) -> T:
+        if durable.is_complete:
+            return durable.get_result()
+        future = self.create_future()
+        entry = self._waits.setdefault(id(durable), (durable, []))
+        entry[1].append(future)
+        if id(durable) not in self._observers:
+            self._observers[id(durable)] = _CompletionObserver(
+                durable, lambda: self._completions.append(lambda: self._complete_wait(durable)))
+        try:
+            return await future
+        finally:
+            # Cancelling one waiter never cancels a shared durable operation.
+            entry[1].remove(future)
+            if not entry[1]:
+                self._waits.pop(id(durable), None)
+                self._observers.pop(id(durable)).detach()
+
+    def _complete_wait(self, durable: task.Task[Any]) -> None:
+        entry = self._waits.get(id(durable))
+        if entry is not None:
+            for future in tuple(entry[1]):
+                if not future.done():
+                    try:
+                        future.set_result(durable.get_result())
+                    except Exception as error:
+                        future.set_exception(error)
+
+    def _deliver_completions(self) -> None:
+        while self._completions:
+            self._completions.popleft()()
+
+    def start(self, coro: Coroutine[Any, Any, Any]) -> None:
+        self._root = self.create_task(coro, name="durable-orchestrator")
+        self.resume()
+
+    def resume(self) -> None:
+        if self._is_complete():
+            return
+        self._deliver_completions()
+        self._drain()
+        assert self._root is not None
+        if self._root.done():
+            if self._root.cancelled():
+                raise task.TaskCancelledError(
+                    "The orchestrator coroutine was cancelled. Catch local cancellation "
+                    "inside the orchestrator; terminate the durable instance through the client.")
+            # result() propagates the orchestrator's failure to the executor.
+            result = self._root.result()
+            raise StopIteration(result)
+        if not self._has_durable_work():
+            raise RuntimeError(
+                "The async orchestrator is blocked without durable work. "
+                "Await durable activities, events or timers; external I/O is unsupported.")
+
+    def _has_durable_work(self) -> bool:
+        return any(not durable.is_complete for durable, _ in self._waits.values())
+
+    def _drain(self) -> None:
+        previous_loop = asyncio._get_running_loop()
+        previous_task = asyncio.current_task(previous_loop) if previous_loop is not None else None
+        if previous_task is not None and previous_loop is not None:
+            asyncio.tasks._leave_task(previous_loop, previous_task)
+        asyncio._set_running_loop(self)
+        try:
+            while self._ready:
+                handle = self._ready.popleft()
+                if not handle.cancelled():
+                    handle._run()
+                if self._fatal is not None:
+                    error, self._fatal = self._fatal, None
+                    if not self._closing:
+                        raise error
+                    self._logger.warning(
+                        "%s: Async cleanup callback failed.", self.ctx.instance_id, exc_info=error)
+                if not self._closing:
+                    if self._is_complete():
+                        break
+                    if self._root is not None and self._root.done():
+                        break
+                    self._deliver_completions()
+        finally:
+            asyncio._set_running_loop(previous_loop)
+            if previous_task is not None and previous_loop is not None:
+                asyncio.tasks._enter_task(previous_loop, previous_task)
+
+    def call_exception_handler(self, context: dict[str, Any]) -> None:
+        error = context.get("exception")
+        if isinstance(error, StopIteration):
+            failure = RuntimeError("An orchestration callback raised StopIteration.")
+            failure.__cause__ = error
+            error = failure
+        self._fatal = error if isinstance(error, BaseException) else RuntimeError(context["message"])
+
+    def close(self) -> None:
+        """Cancel unfinished coroutines so their cleanup runs before the dispatch ends."""
+        if self._closed:
+            return
+        self._closing = True
+        try:
+            for future in tuple(self._tasks):
+                if not future.done():
+                    future.cancel()
+            self._drain()
+            failure: BaseException | None = None
+            for future in tuple(self._tasks):
+                if not future.done():
+                    future.get_coro().close()
+                    failure = failure or RuntimeError(
+                        "Async cleanup must finish without waiting for durable or external work.")
+                elif not future.cancelled() and getattr(future, "_log_traceback", False):
+                    self._logger.warning(
+                        "%s: Async background task failed.", self.ctx.instance_id,
+                        exc_info=future.exception())
+            if failure is not None:
+                raise failure
+        finally:
+            self._closed = True
+            self._ready.clear()
+            for observer in self._observers.values():
+                observer.detach()
+            self._completions.clear()
+            self._observers.clear()
+            self._waits.clear()
+            self._tasks.clear()
diff --git a/durabletask/internal/shared.py b/durabletask/internal/shared.py
index dc000532..079cdfaa 100644
--- a/durabletask/internal/shared.py
+++ b/durabletask/internal/shared.py
@@ -5,6 +5,7 @@
 import logging
 import warnings
 from collections.abc import Sequence
+from functools import partial
 from typing import Any, TypeAlias
 
 import grpc
@@ -73,6 +74,23 @@ def from_json(json_str: str | bytes | bytearray, expected_type: type | None = No
 INSECURE_PROTOCOLS = ["http://", "grpc://"]
 
 
+def is_async_callable(fn: Any) -> bool:
+    """Recognize async functions, wrapped functions, partials and callable objects."""
+    while True:
+        if inspect.iscoroutinefunction(fn):
+            return True
+        fn = inspect.unwrap(fn)
+        if isinstance(fn, partial):
+            fn = fn.func
+            continue
+        if inspect.iscoroutinefunction(fn):
+            return True
+        if inspect.isfunction(fn) or inspect.ismethod(fn) or inspect.isbuiltin(fn):
+            return False
+        call = getattr(fn, "__call__", None)
+        return call is not None and inspect.iscoroutinefunction(inspect.unwrap(call))
+
+
 def _get_legacy_logging_warning_stacklevel() -> int:
     stacklevel = 1
     frame = inspect.currentframe()
diff --git a/durabletask/task.py b/durabletask/task.py
index df2e03a2..79132ea9 100644
--- a/durabletask/task.py
+++ b/durabletask/task.py
@@ -6,8 +6,9 @@
 
 import logging
 import math
+import asyncio
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Generator, Sequence
+from collections.abc import Callable, Coroutine, Generator, Sequence
 from dataclasses import dataclass
 from datetime import UTC, datetime, timedelta
 from typing import Any, Generic, TypeAlias, TypeVar, cast, overload
@@ -612,6 +613,31 @@ def __init__(self) -> None:
         # Dict insertion order preserves subscription order for shared races.
         self._parents = {}
 
+    def __await__(self) -> Generator[Task[T], Any, T]:
+        """Wait for this task in a replay-driven ``async def`` orchestrator.
+
+        A local asyncio waiter does not own or cancel the durable operation.
+        """
+        from durabletask.internal.asyncio_loop import DurableEventLoop
+
+        loop = asyncio._get_running_loop()
+        if isinstance(loop, DurableEventLoop):
+            return (yield from loop.wait_durable(self).__await__())
+        if loop is not None:
+            raise RuntimeError("Durable tasks can only be awaited inside a durable orchestrator.")
+        # Allow mock-based orchestration drivers to send results without an event loop.
+        if not self.is_complete:
+            yield self
+        return self.get_result()
+
+    async def wait(self) -> T:
+        """Borrow this operation's result through a fresh native coroutine.
+
+        Calling this method does not schedule another durable operation. Each
+        call creates an independent waiter.
+        """
+        return await self
+
     @property
     def is_complete(self) -> bool:
         """Returns True if the task has completed, False otherwise."""
@@ -951,8 +977,17 @@ def task_id(self) -> int:
         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]
+# Generator and immediate-result orchestrators use eager durable task handles.
+SyncOrchestrator: TypeAlias = Callable[
+    [OrchestrationContext, TInput],
+    Generator[Task[Any], Any, TOutput] | TOutput,
+]
+
+AsyncOrchestrator: TypeAlias = Callable[
+    [OrchestrationContext, TInput], Coroutine[Any, Any, TOutput],
+]
+
+Orchestrator: TypeAlias = SyncOrchestrator[TInput, TOutput] | AsyncOrchestrator[TInput, TOutput]
 
 # Activities are simple functions that can be scheduled by orchestrators
 Activity: TypeAlias = Callable[[ActivityContext, TInput], TOutput]
diff --git a/durabletask/worker.py b/durabletask/worker.py
index 4c638f14..fa1c5f67 100644
--- a/durabletask/worker.py
+++ b/durabletask/worker.py
@@ -8,12 +8,12 @@
 import logging
 import os
 import time
-from collections.abc import Callable, Generator, Sequence
+from collections.abc import Callable, Coroutine, Generator, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from datetime import UTC, datetime, timedelta
 from threading import Event, Lock, Thread
-from types import GeneratorType
+from types import CoroutineType, GeneratorType
 from enum import Enum
 from functools import partial
 from typing import Any, Self, TypeVar, cast, overload
@@ -50,6 +50,7 @@
 from durabletask.payload import helpers as payload_helpers
 from durabletask import task
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
+from durabletask.internal.asyncio_loop import DurableEventLoop
 from durabletask.payload.store import PayloadStore
 from durabletask.serialization import DataConverter, JsonDataConverter
 
@@ -1669,22 +1670,42 @@ def __init__(self,
         self._data_converter = data_converter
         self._exception_properties_provider = exception_properties_provider
         self._logger = logger
+        self._async_loop: DurableEventLoop | None = None
 
-    def run(self, generator: Generator[task.Task[Any], Any, Any]) -> None:
-        self._generator = generator
-        # TODO: Do something with this task
-        task = next(generator)  # this starts the generator
-        # TODO: Check if the task is null?
-        self._previous_task = task
+    def run(self, orchestrator: Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any]) -> None:
+        if isinstance(orchestrator, Coroutine):
+            self._async_loop = DurableEventLoop(
+                self, is_complete=lambda: self._is_complete, logger=self._logger)
+            self._async_loop.start(orchestrator)
+            return
+        self._generator = orchestrator
+        self._set_next_task(orchestrator.send(None))
+
+    def close(self) -> None:
+        if self._async_loop is not None:
+            self._async_loop.close()
+        elif self._generator is not None:
+            self._generator.close()
+
+    def _set_next_task(self, next_task: Any) -> None:
+        if not isinstance(next_task, task.Task):
+            raise TypeError(
+                "Generator orchestrators must yield durabletask.Task objects. "
+                "Use async def and await for asyncio composition."
+            )
+        self._previous_task = next_task
 
     def resume(self) -> None:
+        if self._async_loop is not None:
+            self._async_loop.resume()
+            return
         if self._generator is None:
             # This is never expected unless maybe there's an issue with the history
             raise TypeError(
-                "The orchestrator generator is not initialized! Was the orchestration history corrupted?"
+                "The orchestrator is not initialized! Was the orchestration history corrupted?"
             )
 
-        # We can resume the generator only if the previously yielded task
+        # We can resume the generator only if the previous task
         # has reached a completed state. The only time this won't be the
         # case is if the user yielded on a WhenAll task and there are still
         # outstanding child tasks that need to be completed.
@@ -1699,9 +1720,7 @@ def resume(self) -> None:
                 # This will either return a Task or raise StopIteration if it's done.
                 next_task = self._generator.send(self._previous_task.get_result())
 
-            if not isinstance(next_task, task.Task):
-                raise TypeError("The orchestrator generator yielded a non-Task object")
-            self._previous_task = next_task
+            self._set_next_task(next_task)
 
     def set_complete(
             self,
@@ -2524,34 +2543,33 @@ def execute(
         finally:
             # Close after building the response, so the orchestrator's finally
             # blocks do not change the actions or custom status we return.
-            if ctx._generator is not None:  # pyright: ignore[reportPrivateUsage]
-                try:
-                    ctx._generator.close()  # pyright: ignore[reportPrivateUsage]
-                except (Exception, asyncio.CancelledError):
-                    cleanup_name = next(
-                        (e.executionStarted.name for e in new_events if e.HasField("executionStarted")),
-                        orchestration_name,
-                    )
-                    self._logger.warning(
-                        f"{instance_id}: Failed to close orchestrator {cleanup_name} generator.",
-                        exc_info=True,
-                    )
-                except BaseException as ex:
-                    self._logger.error(
-                        f"{instance_id}: Orchestrator {orchestration_name} generator teardown exited unexpectedly.",
-                        exc_info=True,
-                    )
-                    if not dispatch_exited:
-                        if _is_control_flow_exception(ex) or result is None:
-                            raise
-                        # Keep an existing failure rather than mask its cause.
-                        if not any(
-                            action.HasField("completeOrchestration")
-                            and action.completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
-                            for action in result.actions
-                        ):
-                            result = result.as_failed(ph.new_failure_details(
-                                ex, self._exception_properties_provider, self._logger))
+            try:
+                ctx.close()
+            except (Exception, asyncio.CancelledError):
+                cleanup_name = next(
+                    (e.executionStarted.name for e in new_events if e.HasField("executionStarted")),
+                    orchestration_name,
+                )
+                self._logger.warning(
+                    f"{instance_id}: Failed to close orchestrator {cleanup_name}.",
+                    exc_info=True,
+                )
+            except BaseException as ex:
+                self._logger.error(
+                    f"{instance_id}: Orchestrator {orchestration_name} teardown exited unexpectedly.",
+                    exc_info=True,
+                )
+                if not dispatch_exited:
+                    if _is_control_flow_exception(ex) or result is None:
+                        raise
+                    # Keep an existing failure rather than mask its cause.
+                    if not any(
+                        action.HasField("completeOrchestration")
+                        and action.completeOrchestration.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+                        for action in result.actions
+                    ):
+                        result = result.as_failed(ph.new_failure_details(
+                            ex, self._exception_properties_provider, self._logger))
         assert result is not None
         return result
 
@@ -2726,12 +2744,28 @@ def process_event(
                     input = self._data_converter.deserialize(
                         event.executionStarted.input.value, input_type)
 
-                result = fn(
-                    ctx, input
-                )  # this does not execute the generator, only creates it
-                if isinstance(result, GeneratorType):
-                    # Start the orchestrator's generator function
-                    ctx.run(cast(Generator[task.Task[Any], Any, Any], result))
+                # Layer 4 passes async orchestrators a cold context instead.
+                is_async = shared.is_async_callable(fn)
+                result = fn(ctx, input)
+
+                if isinstance(result, CoroutineType) and not is_async:
+                    result.close()
+                    raise TypeError(
+                        "A synchronous orchestrator returned a coroutine. Declare it with "
+                        "'async def' or use functools.wraps when decorating an async orchestrator."
+                    )
+                if is_async and not isinstance(result, CoroutineType):
+                    if isinstance(result, GeneratorType):
+                        result.close()
+                    raise TypeError("An async orchestrator must return a coroutine.")
+                if isinstance(result, (GeneratorType, CoroutineType)):
+                    # Both protocols suspend on durable tasks and finish with StopIteration.
+                    ctx.run(cast(Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any], result))
+                elif inspect.isasyncgen(result):
+                    raise TypeError(
+                        "Async generator orchestrators are not supported. "
+                        "Use async def with await, not yield."
+                    )
                 else:
                     # This is an orchestrator that doesn't schedule any tasks
                     ctx.set_complete(result, pb.ORCHESTRATION_STATUS_COMPLETED)
diff --git a/tests/durabletask/test_coroutine_orchestrators.py b/tests/durabletask/test_coroutine_orchestrators.py
new file mode 100644
index 00000000..ee2ca008
--- /dev/null
+++ b/tests/durabletask/test_coroutine_orchestrators.py
@@ -0,0 +1,445 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Replay-level coverage for coroutine orchestrators and generator history parity."""
+
+import asyncio
+import inspect
+import json
+import logging
+from collections.abc import AsyncGenerator, Generator
+from datetime import datetime, timedelta
+from functools import partial, wraps
+from typing import Any
+
+import pytest
+
+from durabletask import entities, task, worker
+from durabletask.internal import helpers
+from durabletask.internal import orchestrator_service_pb2 as pb
+from durabletask.serialization import JsonDataConverter
+
+
+INSTANCE_ID = "coroutine-test"
+NOW = datetime(2026, 1, 1)
+
+
+def _start() -> list[pb.HistoryEvent]:
+    return [
+        helpers.new_orchestrator_started_event(NOW),
+        helpers.new_execution_started_event("test", INSTANCE_ID, json.dumps(3)),
+    ]
+
+
+def _execute(
+    fn: task.Orchestrator[Any, Any],
+    old: list[pb.HistoryEvent],
+    new: list[pb.HistoryEvent],
+) -> worker.ExecutionResults:
+    registry = worker._Registry()  # pyright: ignore[reportPrivateUsage]
+    registry.add_named_orchestrator("test", fn)
+    executor = worker._OrchestrationExecutor(  # pyright: ignore[reportPrivateUsage]
+        registry, logging.getLogger(__name__), JsonDataConverter())
+    return executor.execute(INSTANCE_ID, old, new)
+
+
+def _completion(result: worker.ExecutionResults) -> pb.CompleteOrchestrationAction:
+    completions = [a.completeOrchestration for a in result.actions if a.HasField("completeOrchestration")]
+    assert len(completions) == 1
+    return completions[0]
+
+
+def test_coroutine_immediate_return() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
+        return value
+
+    complete = _completion(_execute(orchestrator, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert json.loads(complete.result.value) == 3
+
+
+@pytest.mark.parametrize("shape", ["function", "wrapped", "partial", "object", "partial-object"])
+def test_async_orchestrator_is_dispatched_for_supported_callable_shapes(shape: str) -> None:
+    async def plain(ctx: task.OrchestrationContext, value: int) -> int:
+        return value + 1
+
+    @wraps(plain)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+        return plain(*args, **kwargs)
+
+    async def with_prefix(prefix: int, ctx: task.OrchestrationContext, value: int) -> int:
+        return value + prefix
+
+    class AsyncCallable:
+        async def __call__(self, ctx: task.OrchestrationContext, value: int) -> int:
+            return value + 1
+
+    class AsyncPartialCallable:
+        async def __call__(
+            self, prefix: int, ctx: task.OrchestrationContext, value: int,
+        ) -> int:
+            return value + prefix
+
+    variants = {
+        "function": plain,
+        "wrapped": wrapped,
+        "partial": partial(with_prefix, 1),
+        "object": AsyncCallable(),
+        "partial-object": partial(AsyncPartialCallable(), 1),
+    }
+    complete = _completion(_execute(variants[shape], [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert json.loads(complete.result.value) == 4
+
+
+def test_unmarked_synchronous_wrapper_returning_coroutine_is_rejected_and_closed() -> None:
+    created: list[Any] = []
+
+    async def async_orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
+        return value
+
+    def wrapper(ctx: task.OrchestrationContext, value: int) -> Any:
+        coroutine = async_orchestrator(ctx, value)  # type: ignore[arg-type]
+        created.append(coroutine)
+        return coroutine
+
+    complete = _completion(_execute(wrapper, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert "synchronous orchestrator returned a coroutine" in complete.failureDetails.errorMessage
+    assert inspect.getcoroutinestate(created[0]) == inspect.CORO_CLOSED
+
+
+def test_async_classified_wrapper_returning_non_coroutine_is_rejected() -> None:
+    async def async_orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
+        return value
+
+    @wraps(async_orchestrator)
+    def wrapper(*args: Any, **kwargs: Any) -> int:
+        return 42
+
+    complete = _completion(_execute(wrapper, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert "async orchestrator must return a coroutine" in complete.failureDetails.errorMessage
+
+
+def test_async_classified_wrapper_closes_returned_generator() -> None:
+    created: list[Generator[int, None, None]] = []
+
+    async def async_orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
+        return value
+
+    @wraps(async_orchestrator)
+    def wrapper(*args: Any, **kwargs: Any) -> Generator[int, None, None]:
+        generator = (value for value in [42])
+        created.append(generator)
+        return generator
+
+    complete = _completion(_execute(wrapper, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert "async orchestrator must return a coroutine" in complete.failureDetails.errorMessage
+    assert inspect.getgeneratorstate(created[0]) == inspect.GEN_CLOSED
+
+
+def test_cyclic_async_wrapper_fails_instead_of_hanging() -> None:
+    def wrapper(ctx: task.OrchestrationContext, value: int) -> int:
+        return value
+
+    wrapper.__wrapped__ = wrapper  # type: ignore[attr-defined]
+    complete = _completion(_execute(wrapper, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert "wrapper loop" in complete.failureDetails.errorMessage
+
+
+def test_await_is_concrete_and_completed_tasks_do_not_suspend() -> None:
+    assert "__await__" not in task.Task.__abstractmethods__
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[int]:
+        completed = task.CompletableTask[int]()
+        completed.complete(7)
+        return [await completed, await completed]
+
+    assert json.loads(_completion(_execute(orchestrator, [], _start())).result.value) == [7, 7]
+
+
+def test_task_wait_suspends_until_durable_completion() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str:
+        return await ctx.call_activity("work", return_type=str).wait()
+
+    old = _start()
+    actions = _execute(orchestrator, [], old).actions
+    assert [a.scheduleTask.name for a in actions] == ["work"]
+    old += [helpers.new_task_scheduled_event(1, "work")]
+    complete = _completion(_execute(
+        orchestrator, old, [helpers.new_task_completed_event(1, '"done"')]))
+    assert json.loads(complete.result.value) == "done"
+
+
+@pytest.mark.parametrize("after_activity", [False, True])
+def test_generator_invalid_yield_recommends_async_authoring(after_activity: bool) -> None:
+    def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[Any, Any, None]:
+        if after_activity:
+            yield ctx.call_activity("work")
+        yield None
+
+    old = _start() + [helpers.new_task_scheduled_event(1, "work")] if after_activity else []
+    new = [helpers.new_task_completed_event(1, "null")] if after_activity else _start()
+    complete = _completion(_execute(orchestrator, old, new))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert complete.failureDetails.errorType == "builtins.TypeError"
+    assert "Generator orchestrators must yield durabletask.Task objects" in complete.failureDetails.errorMessage
+    assert "Use async def and await for asyncio composition" in complete.failureDetails.errorMessage
+
+
+def test_precompleted_failure_and_cancellation_are_catchable() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]:
+        failed = task.CompletableTask[int]()
+        failed.fail("failed", ValueError("boom"))
+        cancelled = task.TimerTask()
+        cancelled.cancel()
+        results: list[str] = []
+        try:
+            await failed
+        except task.TaskFailedError:
+            results.append("failed")
+        try:
+            await cancelled
+        except task.TaskCancelledError:
+            results.append("cancelled")
+        return results
+
+    assert json.loads(_completion(_execute(orchestrator, [], _start())).result.value) == [
+        "failed", "cancelled"]
+
+
+def test_nested_async_helpers_replay_with_generator_action_parity() -> None:
+    async def helper(ctx: task.OrchestrationContext, value: int) -> int:
+        return await ctx.call_activity("increment", input=value, return_type=int)
+
+    async def coroutine(ctx: task.OrchestrationContext, value: int) -> int:
+        first = await helper(ctx, value)
+        return await helper(ctx, first)
+
+    def generator(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, int]:
+        first = yield ctx.call_activity("increment", input=value, return_type=int)
+        return (yield ctx.call_activity("increment", input=first, return_type=int))
+
+    batches = [
+        _start(),
+        [helpers.new_task_scheduled_event(1, "increment"),
+         helpers.new_task_completed_event(1, json.dumps(4))],
+        [helpers.new_task_scheduled_event(2, "increment"),
+         helpers.new_task_completed_event(2, json.dumps(5))],
+    ]
+    old: list[pb.HistoryEvent] = []
+    result: worker.ExecutionResults | None = None
+    for new in batches:
+        result = _execute(coroutine, old, new)
+        assert result.actions == _execute(generator, old, new).actions
+        old += new
+    assert result is not None
+    assert json.loads(_completion(result).result.value) == 5
+
+
+def test_fan_out_waits_for_all_and_preserves_input_order() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]:
+        pending: list[task.Task[str]] = [
+            ctx.call_activity("work", input=i, return_type=str) for i in range(2)]
+        return await task.when_all(pending)
+
+    old = _start()
+    assert len(_execute(orchestrator, [], old).actions) == 2
+    old += [helpers.new_task_scheduled_event(i, "work") for i in (1, 2)]
+    second = [helpers.new_task_completed_event(2, '"second"')]
+    assert _execute(orchestrator, old, second).actions == []
+    complete = _completion(_execute(
+        orchestrator, old + second, [helpers.new_task_completed_event(1, '"first"')]))
+    assert json.loads(complete.result.value) == ["first", "second"]
+
+
+@pytest.mark.parametrize("caught", [False, True])
+def test_activity_failure_is_thrown_into_coroutine(caught: bool) -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str:
+        try:
+            await ctx.call_activity("boom")
+        except task.TaskFailedError:
+            if not caught:
+                raise
+            return "caught"
+        return "unexpected"
+
+    old = _start() + [helpers.new_task_scheduled_event(1, "boom")]
+    complete = _completion(_execute(
+        orchestrator, old, [helpers.new_task_failed_event(1, ValueError("boom"))]))
+    assert complete.orchestrationStatus == (
+        pb.ORCHESTRATION_STATUS_COMPLETED if caught else pb.ORCHESTRATION_STATUS_FAILED)
+    if caught:
+        assert json.loads(complete.result.value) == "caught"
+    else:
+        assert complete.failureDetails.innerFailure.errorMessage == "boom"
+
+
+@pytest.mark.parametrize("event_wins", [False, True])
+def test_when_any_timer_and_external_event(event_wins: bool) -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
+        event = ctx.wait_for_external_event("approval")
+        timer = ctx.create_timer(timedelta(seconds=1))
+        winner = await task.when_any([event, timer])
+        if winner is event:
+            timer.cancel()
+            return await event
+        return "timeout"
+
+    deadline = NOW + timedelta(seconds=1)
+    old = _start() + [helpers.new_timer_created_event(1, deadline)]
+    new = [helpers.new_event_raised_event("approval", '"approved"')] if event_wins else [
+        helpers.new_timer_fired_event(1, deadline)]
+    complete = _completion(_execute(orchestrator, old, new))
+    assert json.loads(complete.result.value) == ("approved" if event_wins else "timeout")
+
+
+def test_sub_orchestration_and_replay_flags() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]:
+        before = ctx.is_replaying
+        value = await ctx.call_sub_orchestrator("child", instance_id="child-test", return_type=str)
+        return [value, before, ctx.is_replaying]
+
+    old = _start() + [helpers.new_sub_orchestration_created_event(1, "child", "child-test")]
+    complete = _completion(_execute(
+        orchestrator, old, [helpers.new_sub_orchestration_completed_event(1, '"done"')]))
+    assert json.loads(complete.result.value) == ["done", True, False]
+
+
+def test_retry_timer_resumes_same_coroutine() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
+        return await ctx.call_activity(
+            "retry", retry_policy=task.RetryPolicy(
+                first_retry_interval=timedelta(seconds=1), max_number_of_attempts=2))
+
+    old = _start() + [helpers.new_task_scheduled_event(1, "retry")]
+    failure = [helpers.new_task_failed_event(1, ValueError("retry me"))]
+    retry = _execute(orchestrator, old, failure).actions
+    assert len(retry) == 1
+    assert retry[0].id == 2
+    deadline = retry[0].createTimer.fireAt.ToDatetime()
+    old += failure + [helpers.new_timer_created_event(2, deadline)]
+    fired = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(2, deadline)]
+    scheduled = _execute(orchestrator, old, fired).actions
+    assert len(scheduled) == 1
+    assert scheduled[0].id == 1
+    assert scheduled[0].scheduleTask.name == "retry"
+    old += fired + [helpers.new_task_scheduled_event(1, "retry")]
+    complete = _completion(_execute(
+        orchestrator, old, [helpers.new_task_completed_event(1, '"success"')]))
+    assert json.loads(complete.result.value) == "success"
+
+
+@pytest.mark.parametrize("sub_orchestration", [False, True])
+@pytest.mark.parametrize("replayed", [False, True])
+@pytest.mark.parametrize("caught", [False, True])
+def test_coroutine_non_retriable_failure(
+        sub_orchestration: bool, replayed: bool, caught: bool) -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
+        retry = task.RetryPolicy(first_retry_interval=timedelta(seconds=1), max_number_of_attempts=3)
+        try:
+            if sub_orchestration:
+                return await ctx.call_sub_orchestrator("child", retry_policy=retry)
+            return await ctx.call_activity("child", retry_policy=retry)
+        except task.TaskFailedError as error:
+            if not caught:
+                raise
+            assert error.details.message == "permanent"
+            return "caught"
+
+    old = _start()
+    if sub_orchestration:
+        old.append(helpers.new_sub_orchestration_created_event(1, "child", "child-instance"))
+        failed = helpers.new_sub_orchestration_failed_event(1, ValueError("permanent"))
+        details = failed.subOrchestrationInstanceFailed.failureDetails
+    else:
+        old.append(helpers.new_task_scheduled_event(1, "child"))
+        failed = helpers.new_task_failed_event(1, ValueError("permanent"))
+        details = failed.taskFailed.failureDetails
+    details.isNonRetriable = True
+    new = [helpers.new_orchestrator_started_event(NOW)]
+    if replayed:
+        old.append(failed)
+    else:
+        new.append(failed)
+    result = _execute(orchestrator, old, new)
+    assert len(result.actions) == 1
+    complete = _completion(result)
+    if caught:
+        assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+        assert json.loads(complete.result.value) == "caught"
+    else:
+        assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+        assert complete.failureDetails.innerFailure == details
+
+
+def test_entity_lock_and_call_suspend_coroutine() -> None:
+    entity = entities.EntityInstanceId("counter", "one")
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
+        with await ctx.lock_entities([entity]):
+            return await ctx.call_entity(entity, "get")
+
+    old = _start()
+    lock = _execute(orchestrator, [], old).actions[0].sendEntityMessage.entityLockRequested
+    request = json.dumps({"id": lock.criticalSectionId})
+    new = [
+        helpers.new_event_sent_event(1, str(entity), request),
+        helpers.new_event_raised_event(lock.criticalSectionId, None),
+    ]
+    actions = _execute(orchestrator, old, new).actions
+    assert len(actions) == 1
+    call = actions[0].sendEntityMessage.entityOperationCalled
+    assert call.targetInstanceId.value == str(entity)
+    old += new + [helpers.new_event_sent_event(2, str(entity), json.dumps({"id": call.requestId}))]
+    complete = _completion(_execute(
+        orchestrator, old, [helpers.new_event_raised_event(call.requestId, json.dumps({"result": "42"}))]))
+    assert json.loads(complete.result.value) == 42
+
+
+def test_continue_as_new_from_coroutine() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, value: int) -> None:
+        await ctx.wait_for_external_event("next")
+        ctx.continue_as_new(value + 1)
+
+    complete = _completion(_execute(
+        orchestrator, _start(), [helpers.new_event_raised_event("next", None)]))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW
+    assert json.loads(complete.result.value) == 4
+
+
+def test_nondeterminism_still_detected_for_coroutines() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
+        return await ctx.call_activity("changed")
+
+    complete = _completion(_execute(
+        orchestrator, _start(), [helpers.new_task_scheduled_event(1, "original")]))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert complete.failureDetails.errorType.endswith("NonDeterminismError")
+
+
+@pytest.mark.parametrize("after_activity", [False, True])
+def test_non_durable_suspension_is_rejected(after_activity: bool) -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        if after_activity:
+            await ctx.call_activity("work")
+        await asyncio.Future()
+
+    old = _start() + [helpers.new_task_scheduled_event(1, "work")] if after_activity else []
+    new = [helpers.new_task_completed_event(1, "null")] if after_activity else _start()
+    complete = _completion(_execute(orchestrator, old, new))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert complete.failureDetails.errorType == "builtins.RuntimeError"
+    assert "blocked without durable work" in complete.failureDetails.errorMessage
+
+
+def test_async_generator_is_rejected() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> AsyncGenerator[task.Task[Any], None]:
+        yield ctx.call_activity("work", return_type=int)
+
+    complete = _completion(_execute(orchestrator, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert "Async generator orchestrators are not supported" in complete.failureDetails.errorMessage
diff --git a/tests/durabletask/test_orchestration_generator_cleanup.py b/tests/durabletask/test_orchestration_generator_cleanup.py
index 80adfd3e..4b348d38 100644
--- a/tests/durabletask/test_orchestration_generator_cleanup.py
+++ b/tests/durabletask/test_orchestration_generator_cleanup.py
@@ -347,7 +347,7 @@ def test_teardown_cancellation_preserves_response(
     assert len(caplog.records) == 1
     record = caplog.records[0]
     assert record.levelno == logging.WARNING
-    assert record.message == f"{TEST_INSTANCE_ID}: Failed to close orchestrator orchestrator generator."
+    assert record.message == f"{TEST_INSTANCE_ID}: Failed to close orchestrator orchestrator."
     assert record.exc_info is not None
     assert record.exc_info[0] is asyncio.CancelledError
 

From e73a705eee9e59156ce4b94b30435935aa6b8d1e Mon Sep 17 00:00:00 2001
From: Andy Staples <andystaples@microsoft.com>
Date: Thu, 12 Nov 2026 13:22:19 -0700
Subject: [PATCH 2/3] Fix async callable wrapper traversal

Track identities across wrapped, partial, and callable nodes; recognize intermediate coroutine functions before following their wrapped targets. Cover bounded worker failures and nested decorated dispatch.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 CHANGELOG.md                                  |   4 +
 durabletask/internal/shared.py                |  20 ++-
 .../test_coroutine_orchestrators.py           | 116 +++++++++++++++++-
 3 files changed, 132 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 40aa2e88..15f56206 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,10 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
+- Async orchestrators wrapped with `functools.wraps` are now recognized even
+when an intermediate async wrapper decorates a synchronous function, including
+callable instances and partials. Cyclic wrapper chains fail instead of hanging
+worker dispatch.
 - Activity and sub-orchestration `retry_timeout` decisions now use replay-safe
 orchestration time, preventing retries from expiring solely because history is
 replayed later. Retry calculations also accept timezone-aware timestamps.
diff --git a/durabletask/internal/shared.py b/durabletask/internal/shared.py
index 079cdfaa..26541738 100644
--- a/durabletask/internal/shared.py
+++ b/durabletask/internal/shared.py
@@ -76,19 +76,27 @@ def from_json(json_str: str | bytes | bytearray, expected_type: type | None = No
 
 def is_async_callable(fn: Any) -> bool:
     """Recognize async functions, wrapped functions, partials and callable objects."""
+    # Keep nodes alive so temporary bound-method IDs cannot be reused.
+    seen: dict[int, Any] = {}
     while True:
+        identity = id(fn)
+        if identity in seen:
+            raise ValueError("wrapper loop when classifying async callable")
+        seen[identity] = fn
         if inspect.iscoroutinefunction(fn):
             return True
-        fn = inspect.unwrap(fn)
+        if hasattr(fn, "__wrapped__"):
+            fn = fn.__wrapped__
+            continue
         if isinstance(fn, partial):
             fn = fn.func
             continue
-        if inspect.iscoroutinefunction(fn):
-            return True
-        if inspect.isfunction(fn) or inspect.ismethod(fn) or inspect.isbuiltin(fn):
+        if (inspect.isfunction(fn) or inspect.ismethod(fn)
+                or inspect.isbuiltin(fn) or inspect.ismethodwrapper(fn)):
+            return False
+        fn = getattr(fn, "__call__", None)
+        if fn is None:
             return False
-        call = getattr(fn, "__call__", None)
-        return call is not None and inspect.iscoroutinefunction(inspect.unwrap(call))
 
 
 def _get_legacy_logging_warning_stacklevel() -> int:
diff --git a/tests/durabletask/test_coroutine_orchestrators.py b/tests/durabletask/test_coroutine_orchestrators.py
index ee2ca008..824b3afe 100644
--- a/tests/durabletask/test_coroutine_orchestrators.py
+++ b/tests/durabletask/test_coroutine_orchestrators.py
@@ -7,15 +7,17 @@
 import inspect
 import json
 import logging
+import subprocess
+import sys
 from collections.abc import AsyncGenerator, Generator
 from datetime import datetime, timedelta
-from functools import partial, wraps
+from functools import partial, update_wrapper, wraps
 from typing import Any
 
 import pytest
 
 from durabletask import entities, task, worker
-from durabletask.internal import helpers
+from durabletask.internal import helpers, shared
 from durabletask.internal import orchestrator_service_pb2 as pb
 from durabletask.serialization import JsonDataConverter
 
@@ -92,6 +94,75 @@ async def __call__(
     assert json.loads(complete.result.value) == 4
 
 
+@pytest.mark.parametrize("shape", [
+    "function", "partial", "bound-method", "callable", "partial-callable",
+    "async-callable", "nested-callable",
+])
+def test_async_orchestrator_is_recognized_before_its_synchronous_wrapped_function(shape: str) -> None:
+    def base(ctx: task.OrchestrationContext, value: int) -> int:
+        return value + 1
+
+    @wraps(base)
+    async def asynchronous(ctx: task.OrchestrationContext, value: int) -> int:
+        return base(ctx, value)
+
+    @wraps(asynchronous)
+    def traced(ctx: task.OrchestrationContext, value: int) -> Any:
+        return asynchronous(ctx, value)
+
+    class DecoratedCallable:
+        def base(self, ctx: task.OrchestrationContext, value: int) -> int:
+            return value + 1
+
+        async def asynchronous(self, ctx: task.OrchestrationContext, value: int) -> int:
+            return self.base(ctx, value)
+
+        def __call__(self, ctx: task.OrchestrationContext, value: int) -> Any:
+            return self.asynchronous(ctx, value)
+
+    update_wrapper(DecoratedCallable.asynchronous, DecoratedCallable.base)
+    update_wrapper(DecoratedCallable.__call__, DecoratedCallable.asynchronous)
+
+    class AsyncCallable(DecoratedCallable):
+        async def __call__(self, ctx: task.OrchestrationContext, value: int) -> int:
+            return self.base(ctx, value)
+
+    update_wrapper(AsyncCallable.__call__, AsyncCallable.base)
+
+    def wrap_callable(target: task.Orchestrator[Any, Any]) -> task.Orchestrator[Any, Any]:
+        class Wrapper:
+            def __call__(self, ctx: task.OrchestrationContext, value: int) -> Any:
+                return target(ctx, value)
+
+        update_wrapper(Wrapper.__call__, target)
+        return Wrapper()
+
+    decorated = DecoratedCallable()
+    nested: task.Orchestrator[Any, Any] = traced
+    # Each __call__ lookup produces a temporary bound method.
+    for _ in range(32):
+        nested = wrap_callable(nested)
+    variants = {
+        "function": traced,
+        "partial": partial(traced),
+        "bound-method": decorated.__call__,
+        "callable": decorated,
+        "partial-callable": partial(decorated),
+        "async-callable": AsyncCallable(),
+        "nested-callable": nested,
+    }
+    fn = variants[shape]
+    assert shared.is_async_callable(fn)
+    complete = _completion(_execute(fn, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert json.loads(complete.result.value) == 4
+
+
+@pytest.mark.parametrize("fn", [int, len, list[int].append, object().__str__])
+def test_synchronous_builtin_callables_are_not_classified_as_async(fn: Any) -> None:
+    assert not shared.is_async_callable(fn)
+
+
 def test_unmarked_synchronous_wrapper_returning_coroutine_is_rejected_and_closed() -> None:
     created: list[Any] = []
 
@@ -150,6 +221,47 @@ def wrapper(ctx: task.OrchestrationContext, value: int) -> int:
     assert "wrapper loop" in complete.failureDetails.errorMessage
 
 
+def _assert_mixed_wrapper_cycle_fails(shape: str) -> None:
+    def wrapper(ctx: task.OrchestrationContext, value: int) -> int:
+        return value
+
+    if shape == "partial":
+        wrapper.__wrapped__ = partial(wrapper)  # type: ignore[attr-defined]
+    else:
+        class CallableWrapper:
+            def __call__(self, ctx: task.OrchestrationContext, value: int) -> int:
+                return value
+
+        wrapper.__wrapped__ = CallableWrapper()  # type: ignore[attr-defined]
+        CallableWrapper.__call__.__wrapped__ = partial(wrapper)  # type: ignore[attr-defined]
+
+    # No input: dispatch must detect the cycle without relying on type discovery.
+    history = [
+        helpers.new_orchestrator_started_event(NOW),
+        helpers.new_execution_started_event("test", INSTANCE_ID),
+    ]
+    complete = _completion(_execute(wrapper, [], history))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert complete.failureDetails.errorType == "builtins.ValueError"
+    assert "wrapper loop" in complete.failureDetails.errorMessage
+    with pytest.raises(ValueError, match="wrapper loop"):
+        shared.is_async_callable(wrapper)
+
+
+@pytest.mark.parametrize("shape", ["partial", "callable-partial"])
+def test_mixed_wrapper_cycle_fails_within_bounded_worker_dispatch(shape: str) -> None:
+    result = subprocess.run(
+        [
+            sys.executable, "-c",
+            "import runpy, sys; "
+            "runpy.run_path(sys.argv[1])[sys.argv[2]](sys.argv[3])",
+            __file__, _assert_mixed_wrapper_cycle_fails.__name__, shape,
+        ],
+        capture_output=True, text=True, timeout=10,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
 def test_await_is_concrete_and_completed_tasks_do_not_suspend() -> None:
     assert "__await__" not in task.Task.__abstractmethods__
 

From 1acb27651fe27af1c37989999ff27c35744cf2d6 Mon Sep 17 00:00:00 2001
From: Andy Staples <andystaples@microsoft.com>
Date: Thu, 12 Nov 2026 15:45:22 -0700
Subject: [PATCH 3/3] Classify orchestrators at registration and harden the
 coroutine loop

Check each callable before following functools.wraps metadata so generator
wrappers of async functions stay synchronous, and classify once when the
orchestrator is registered so wrapper cycles fail at registration. Normalize
timezone-aware orchestration clocks in loop.time(), give loop timers an
actionable error, document the private asyncio APIs the loop depends on, and
add a generator history replay compatibility test. Remove the changelog entry
for unreleased behavior.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 CHANGELOG.md                                  |   4 -
 durabletask/internal/asyncio_loop.py          |  36 +++++-
 durabletask/internal/shared.py                |  42 ++++---
 durabletask/task.py                           |   2 +-
 durabletask/worker.py                         |  18 ++-
 .../test_coroutine_orchestrators.py           |  90 ++++++++++++---
 .../test_generator_replay_compat.py           | 103 ++++++++++++++++++
 7 files changed, 254 insertions(+), 41 deletions(-)
 create mode 100644 tests/durabletask/test_generator_replay_compat.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 15f56206..40aa2e88 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,10 +9,6 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
-- Async orchestrators wrapped with `functools.wraps` are now recognized even
-when an intermediate async wrapper decorates a synchronous function, including
-callable instances and partials. Cyclic wrapper chains fail instead of hanging
-worker dispatch.
 - Activity and sub-orchestration `retry_timeout` decisions now use replay-safe
 orchestration time, preventing retries from expiring solely because history is
 replayed later. Retry calculations also accept timezone-aware timestamps.
diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py
index 8865123b..ef97c4aa 100644
--- a/durabletask/internal/asyncio_loop.py
+++ b/durabletask/internal/asyncio_loop.py
@@ -1,7 +1,22 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-"""Dispatch-local asyncio scheduling on top of the shared durable history engine."""
+"""Dispatch-local asyncio scheduling on top of the shared durable history engine.
+
+Private asyncio dependencies (verified against the supported Python versions):
+
+- ``asyncio._get_running_loop`` / ``asyncio._set_running_loop``: install this
+  loop as the running loop only while a dispatch drains ready callbacks, and
+  restore any host loop afterwards. ``Task.__await__`` uses the getter to
+  detect durable, foreign, or absent loops without raising.
+- ``asyncio.tasks._enter_task`` / ``asyncio.tasks._leave_task``: suspend and
+  restore the host's current task while draining, so asyncio does not see two
+  tasks running on one thread.
+- ``asyncio.Handle._run``: run a ready callback with asyncio's own context and
+  exception-handler semantics.
+- ``asyncio.Task._log_traceback``: report unobserved local-task failures during
+  dispatch cleanup without marking them as retrieved.
+"""
 
 from __future__ import annotations
 
@@ -10,13 +25,15 @@
 import logging
 from collections import deque
 from collections.abc import Callable, Coroutine, Generator
-from datetime import datetime
+from datetime import UTC, datetime
 from typing import Any, TypeVar
 
 from durabletask import task
 
 T = TypeVar("T")
 
+_EPOCH = datetime(1970, 1, 1)
+
 
 class _CompletionObserver(task.CompositeTask[None]):
     def __init__(self, durable: task.Task[Any], callback: Callable[[], None]) -> None:
@@ -59,7 +76,10 @@ def __init__(
 
     def time(self) -> float:
         # A fixed UTC epoch avoids local-time and platform timestamp differences.
-        return (self.ctx.current_utc_datetime - datetime(1970, 1, 1)).total_seconds()
+        now = self.ctx.current_utc_datetime
+        if now.tzinfo is not None:
+            now = now.astimezone(UTC).replace(tzinfo=None)
+        return (now - _EPOCH).total_seconds()
 
     def get_debug(self) -> bool:
         return False
@@ -98,6 +118,16 @@ def call_soon(
         self._ready.append(handle)
         return handle
 
+    def call_later(self, delay: float, callback: Callable[..., Any], *args: Any,
+                   context: contextvars.Context | None = None) -> asyncio.TimerHandle:
+        raise NotImplementedError(
+            "Loop timers, including asyncio.sleep() with a positive delay, are not "
+            "supported by the durable orchestration loop. Await ctx.create_timer(...) for a durable timer.")
+
+    def call_at(self, when: float, callback: Callable[..., Any], *args: Any,
+                context: contextvars.Context | None = None) -> asyncio.TimerHandle:
+        return self.call_later(when - self.time(), callback, *args, context=context)
+
     async def wait_durable(self, durable: task.Task[T]) -> T:
         if durable.is_complete:
             return durable.get_result()
diff --git a/durabletask/internal/shared.py b/durabletask/internal/shared.py
index 26541738..43e71e41 100644
--- a/durabletask/internal/shared.py
+++ b/durabletask/internal/shared.py
@@ -75,28 +75,40 @@ def from_json(json_str: str | bytes | bytearray, expected_type: type | None = No
 
 
 def is_async_callable(fn: Any) -> bool:
-    """Recognize async functions, wrapped functions, partials and callable objects."""
+    """Classify an orchestrator callable as async (``True``) or synchronous.
+
+    Each callable is checked directly before following ``functools.wraps``
+    metadata: coroutine functions are async and generator functions are
+    synchronous, even when they wrap a function of the other kind. Partials and
+    callable instances are resolved through ``func`` and ``__call__``. Only a
+    callable that is neither falls back to the function it wraps.
+    """
     # Keep nodes alive so temporary bound-method IDs cannot be reused.
     seen: dict[int, Any] = {}
+    wrapped_fallbacks: list[Any] = []
+    node: Any = fn
     while True:
-        identity = id(fn)
+        if node is None:
+            if not wrapped_fallbacks:
+                return False
+            node = wrapped_fallbacks.pop()
+        identity = id(node)
         if identity in seen:
             raise ValueError("wrapper loop when classifying async callable")
-        seen[identity] = fn
-        if inspect.iscoroutinefunction(fn):
+        seen[identity] = node
+        if inspect.iscoroutinefunction(node):
             return True
-        if hasattr(fn, "__wrapped__"):
-            fn = fn.__wrapped__
-            continue
-        if isinstance(fn, partial):
-            fn = fn.func
-            continue
-        if (inspect.isfunction(fn) or inspect.ismethod(fn)
-                or inspect.isbuiltin(fn) or inspect.ismethodwrapper(fn)):
-            return False
-        fn = getattr(fn, "__call__", None)
-        if fn is None:
+        if inspect.isgeneratorfunction(node) or inspect.isasyncgenfunction(node):
             return False
+        if hasattr(node, "__wrapped__"):
+            wrapped_fallbacks.append(node.__wrapped__)
+        if isinstance(node, partial):
+            node = node.func
+        elif (inspect.isfunction(node) or inspect.ismethod(node)
+                or inspect.isbuiltin(node) or inspect.ismethodwrapper(node)):
+            node = None
+        else:
+            node = getattr(node, "__call__", None)
 
 
 def _get_legacy_logging_warning_stacklevel() -> int:
diff --git a/durabletask/task.py b/durabletask/task.py
index 79132ea9..44d8fc0a 100644
--- a/durabletask/task.py
+++ b/durabletask/task.py
@@ -4,9 +4,9 @@
 # See https://peps.python.org/pep-0563/
 from __future__ import annotations
 
+import asyncio
 import logging
 import math
-import asyncio
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Coroutine, Generator, Sequence
 from dataclasses import dataclass
diff --git a/durabletask/worker.py b/durabletask/worker.py
index fa1c5f67..e819ea64 100644
--- a/durabletask/worker.py
+++ b/durabletask/worker.py
@@ -358,6 +358,7 @@ class _Registry:
 
     def __init__(self):
         self.orchestrators = {}
+        self._orchestrator_kinds: dict[str, tuple[task.Orchestrator[Any, Any], bool]] = {}
         self.activities = {}
         self.entities = {}
 
@@ -375,11 +376,20 @@ def add_named_orchestrator(self, name: str, fn: task.Orchestrator[TInput, TOutpu
         if name in self.orchestrators:
             raise ValueError(f"A '{name}' orchestrator already exists.")
 
+        # Classify once so wrapper problems surface at registration, not dispatch.
+        is_async = shared.is_async_callable(fn)
         self.orchestrators[name] = fn
+        self._orchestrator_kinds[name] = (fn, is_async)
 
     def get_orchestrator(self, name: str) -> task.Orchestrator[Any, Any] | None:
         return self.orchestrators.get(name)
 
+    def is_async_orchestrator(self, name: str, fn: task.Orchestrator[Any, Any]) -> bool:
+        kind = self._orchestrator_kinds.get(name)
+        if kind is not None and kind[0] is fn:
+            return kind[1]
+        return shared.is_async_callable(fn)
+
     def add_activity(self, fn: task.Activity[TInput, TOutput]) -> str:
         if fn is None:
             raise ValueError("An activity function argument is required.")
@@ -2745,7 +2755,7 @@ def process_event(
                         event.executionStarted.input.value, input_type)
 
                 # Layer 4 passes async orchestrators a cold context instead.
-                is_async = shared.is_async_callable(fn)
+                is_async = self._registry.is_async_orchestrator(event.executionStarted.name, fn)
                 result = fn(ctx, input)
 
                 if isinstance(result, CoroutineType) and not is_async:
@@ -2757,7 +2767,11 @@ def process_event(
                 if is_async and not isinstance(result, CoroutineType):
                     if isinstance(result, GeneratorType):
                         result.close()
-                    raise TypeError("An async orchestrator must return a coroutine.")
+                    raise TypeError(
+                        "An async orchestrator must return a coroutine. A synchronous wrapper "
+                        "decorated with functools.wraps(async_fn) is classified as async; declare "
+                        "the wrapper with 'async def', or do not copy the async function's metadata."
+                    )
                 if isinstance(result, (GeneratorType, CoroutineType)):
                     # Both protocols suspend on durable tasks and finish with StopIteration.
                     ctx.run(cast(Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any], result))
diff --git a/tests/durabletask/test_coroutine_orchestrators.py b/tests/durabletask/test_coroutine_orchestrators.py
index 824b3afe..04b3ce0a 100644
--- a/tests/durabletask/test_coroutine_orchestrators.py
+++ b/tests/durabletask/test_coroutine_orchestrators.py
@@ -10,17 +10,21 @@
 import subprocess
 import sys
 from collections.abc import AsyncGenerator, Generator
-from datetime import datetime, timedelta
+from datetime import UTC, datetime, timedelta, timezone
 from functools import partial, update_wrapper, wraps
-from typing import Any
+from types import SimpleNamespace
+from typing import Any, cast
 
 import pytest
 
 from durabletask import entities, task, worker
 from durabletask.internal import helpers, shared
 from durabletask.internal import orchestrator_service_pb2 as pb
+from durabletask.internal.asyncio_loop import DurableEventLoop
 from durabletask.serialization import JsonDataConverter
 
+pytestmark = pytest.mark.filterwarnings("error::RuntimeWarning")
+
 
 INSTANCE_ID = "coroutine-test"
 NOW = datetime(2026, 1, 1)
@@ -211,14 +215,55 @@ def wrapper(*args: Any, **kwargs: Any) -> Generator[int, None, None]:
     assert inspect.getgeneratorstate(created[0]) == inspect.GEN_CLOSED
 
 
-def test_cyclic_async_wrapper_fails_instead_of_hanging() -> None:
+@pytest.mark.parametrize("shape", ["function", "partial", "callable"])
+def test_generator_wrapper_of_async_function_is_synchronous(shape: str) -> None:
+    async def async_orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
+        return value
+
+    @wraps(async_orchestrator)
+    def generator(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, int]:
+        return (yield ctx.call_activity("work", input=value, return_type=int))
+
+    class GeneratorCallable:
+        def __call__(self, ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, int]:
+            return (yield ctx.call_activity("work", input=value, return_type=int))
+
+    update_wrapper(GeneratorCallable.__call__, async_orchestrator)
+    variants: dict[str, Any] = {
+        "function": generator,
+        "partial": partial(generator),
+        "callable": GeneratorCallable(),
+    }
+    fn = variants[shape]
+    assert not shared.is_async_callable(fn)
+    actions = _execute(fn, [], _start()).actions
+    assert [a.scheduleTask.name for a in actions] == ["work"]
+
+
+def test_orchestrator_is_classified_once_at_registration(monkeypatch: pytest.MonkeyPatch) -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, value: int) -> int:
+        return await ctx.call_activity("work", input=value, return_type=int)
+
+    registry = worker._Registry()  # pyright: ignore[reportPrivateUsage]
+    registry.add_named_orchestrator("test", orchestrator)
+    calls: list[Any] = []
+    original = shared.is_async_callable
+    monkeypatch.setattr(shared, "is_async_callable", lambda fn: calls.append(fn) or original(fn))
+    executor = worker._OrchestrationExecutor(  # pyright: ignore[reportPrivateUsage]
+        registry, logging.getLogger(__name__), JsonDataConverter())
+    old = _start() + [helpers.new_task_scheduled_event(1, "work")]
+    executor.execute(INSTANCE_ID, [], _start())
+    executor.execute(INSTANCE_ID, old, [helpers.new_task_completed_event(1, "4")])
+    assert calls == []
+
+
+def test_cyclic_async_wrapper_is_rejected_at_registration() -> None:
     def wrapper(ctx: task.OrchestrationContext, value: int) -> int:
         return value
 
     wrapper.__wrapped__ = wrapper  # type: ignore[attr-defined]
-    complete = _completion(_execute(wrapper, [], _start()))
-    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
-    assert "wrapper loop" in complete.failureDetails.errorMessage
+    with pytest.raises(ValueError, match="wrapper loop"):
+        worker._Registry().add_named_orchestrator("test", wrapper)  # pyright: ignore[reportPrivateUsage]
 
 
 def _assert_mixed_wrapper_cycle_fails(shape: str) -> None:
@@ -235,21 +280,16 @@ def __call__(self, ctx: task.OrchestrationContext, value: int) -> int:
         wrapper.__wrapped__ = CallableWrapper()  # type: ignore[attr-defined]
         CallableWrapper.__call__.__wrapped__ = partial(wrapper)  # type: ignore[attr-defined]
 
-    # No input: dispatch must detect the cycle without relying on type discovery.
-    history = [
-        helpers.new_orchestrator_started_event(NOW),
-        helpers.new_execution_started_event("test", INSTANCE_ID),
-    ]
-    complete = _completion(_execute(wrapper, [], history))
-    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
-    assert complete.failureDetails.errorType == "builtins.ValueError"
-    assert "wrapper loop" in complete.failureDetails.errorMessage
+    registry = worker._Registry()  # pyright: ignore[reportPrivateUsage]
+    with pytest.raises(ValueError, match="wrapper loop"):
+        registry.add_named_orchestrator("test", wrapper)
+    assert registry.get_orchestrator("test") is None
     with pytest.raises(ValueError, match="wrapper loop"):
         shared.is_async_callable(wrapper)
 
 
 @pytest.mark.parametrize("shape", ["partial", "callable-partial"])
-def test_mixed_wrapper_cycle_fails_within_bounded_worker_dispatch(shape: str) -> None:
+def test_mixed_wrapper_cycle_is_rejected_within_bounded_registration(shape: str) -> None:
     result = subprocess.run(
         [
             sys.executable, "-c",
@@ -555,3 +595,21 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> AsyncGenerator
     complete = _completion(_execute(orchestrator, [], _start()))
     assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
     assert "Async generator orchestrators are not supported" in complete.failureDetails.errorMessage
+
+
+@pytest.mark.parametrize(("now", "offset"), [
+    (NOW, 0), (NOW.replace(tzinfo=UTC), 0), (NOW.replace(tzinfo=timezone(timedelta(hours=2))), 7200),
+])
+def test_loop_time_accepts_naive_and_aware_orchestration_clocks(now: datetime, offset: int) -> None:
+    ctx = cast(task.OrchestrationContext, SimpleNamespace(current_utc_datetime=now))
+    loop = DurableEventLoop(ctx, is_complete=lambda: False)
+    assert loop.time() == (NOW - datetime(1970, 1, 1)).total_seconds() - offset
+
+
+def test_positive_asyncio_sleep_reports_durable_timer_alternative() -> None:
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        await asyncio.sleep(1)
+
+    complete = _completion(_execute(orchestrator, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert "ctx.create_timer" in complete.failureDetails.errorMessage
diff --git a/tests/durabletask/test_generator_replay_compat.py b/tests/durabletask/test_generator_replay_compat.py
new file mode 100644
index 00000000..18530e10
--- /dev/null
+++ b/tests/durabletask/test_generator_replay_compat.py
@@ -0,0 +1,103 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Generator orchestrators must keep replaying histories recorded by earlier SDK releases."""
+
+import json
+import logging
+from collections.abc import Generator
+from datetime import datetime, timedelta
+from typing import Any
+
+from durabletask import task, worker
+from durabletask.internal import helpers
+from durabletask.internal import orchestrator_service_pb2 as pb
+from durabletask.serialization import JsonDataConverter
+
+INSTANCE_ID = "generator-compat"
+NOW = datetime(2026, 1, 1)
+
+
+def _orchestrator(ctx: task.OrchestrationContext, values: list[int]) -> Generator[task.Task[Any], Any, Any]:
+    doubled = yield task.when_all([ctx.call_activity("double", input=value) for value in values])
+    yield ctx.create_timer(timedelta(minutes=5))
+    approval = ctx.wait_for_external_event("approval")
+    winner = yield task.when_any([approval, ctx.create_timer(timedelta(hours=1))])
+    approved = winner == approval and approval.get_result()
+    child = yield ctx.call_sub_orchestrator("child", input=sum(doubled), instance_id="child-1")
+    return {"doubled": doubled, "approved": approved, "child": child}
+
+
+def _execute(old: list[pb.HistoryEvent], new: list[pb.HistoryEvent]) -> list[pb.OrchestratorAction]:
+    registry = worker._Registry()  # pyright: ignore[reportPrivateUsage]
+    registry.add_named_orchestrator("compat", _orchestrator)
+    executor = worker._OrchestrationExecutor(  # pyright: ignore[reportPrivateUsage]
+        registry, logging.getLogger(__name__), JsonDataConverter())
+    return list(executor.execute(INSTANCE_ID, old, new).actions)
+
+
+def _summary(actions: list[pb.OrchestratorAction]) -> list[tuple[Any, ...]]:
+    result: list[tuple[Any, ...]] = []
+    for action in actions:
+        if action.HasField("scheduleTask"):
+            result.append((action.id, "activity", action.scheduleTask.name, action.scheduleTask.input.value))
+        elif action.HasField("createTimer"):
+            result.append((action.id, "timer", action.createTimer.fireAt.ToDatetime().isoformat()))
+        elif action.HasField("createSubOrchestration"):
+            sub = action.createSubOrchestration
+            result.append((action.id, "sub", sub.name, sub.instanceId, sub.input.value))
+        elif action.HasField("completeOrchestration"):
+            complete = action.completeOrchestration
+            result.append((action.id, "complete", complete.orchestrationStatus, complete.result.value))
+        else:
+            result.append((action.id, action.WhichOneof("orchestratorActionType")))
+    return result
+
+
+def test_generator_history_from_previous_releases_replays_unchanged() -> None:
+    timer_at = NOW + timedelta(minutes=5)
+    deadline = timer_at + timedelta(hours=1)
+    old: list[pb.HistoryEvent] = []
+    steps: list[tuple[list[pb.HistoryEvent], list[tuple[Any, ...]]]] = [
+        ([
+            helpers.new_orchestrator_started_event(NOW),
+            helpers.new_execution_started_event("compat", INSTANCE_ID, json.dumps([1, 2])),
+        ], [
+            (1, "activity", "double", "1"),
+            (2, "activity", "double", "2"),
+        ]),
+        ([
+            helpers.new_orchestrator_started_event(NOW),
+            helpers.new_task_scheduled_event(1, "double"),
+            helpers.new_task_scheduled_event(2, "double"),
+            helpers.new_task_completed_event(2, "4"),
+            helpers.new_task_completed_event(1, "2"),
+        ], [
+            (3, "timer", timer_at.isoformat()),
+        ]),
+        ([
+            helpers.new_orchestrator_started_event(timer_at),
+            helpers.new_timer_created_event(3, timer_at),
+            helpers.new_timer_fired_event(3, timer_at),
+        ], [
+            (4, "timer", deadline.isoformat()),
+        ]),
+        ([
+            helpers.new_orchestrator_started_event(timer_at),
+            helpers.new_timer_created_event(4, deadline),
+            helpers.new_event_raised_event("approval", "true"),
+        ], [
+            (5, "sub", "child", "child-1", "6"),
+        ]),
+        ([
+            helpers.new_orchestrator_started_event(timer_at),
+            helpers.new_sub_orchestration_created_event(5, "child", "child-1", "6"),
+            helpers.new_sub_orchestration_completed_event(5, '"done"'),
+        ], [
+            (6, "complete", pb.ORCHESTRATION_STATUS_COMPLETED,
+             json.dumps({"doubled": [2, 4], "approved": True, "child": "done"})),
+        ]),
+    ]
+    for new, expected in steps:
+        assert _summary(_execute(old, new)) == expected
+        old += new