From fd6b4025b58b2fd12dca956620abf08ad37a6102 Mon Sep 17 00:00:00 2001
From: Andy Staples <andystaples@microsoft.com>
Date: Wed, 11 Nov 2026 12:59:07 -0700
Subject: [PATCH 1/4] Harden async orchestration dispatch cleanup

Close async generators and drain or cancel pending tasks at dispatch teardown, share control-flow exception classification, and add a Windows async-runtime CI job.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .github/workflows/durabletask.yml             |  24 ++
 durabletask/internal/asyncio_loop.py          |  87 ++++-
 durabletask/internal/exceptions.py            |  11 +
 durabletask/worker.py                         |  18 +-
 .../durabletask/test_asyncio_orchestrators.py | 299 ++++++++++++++++++
 5 files changed, 411 insertions(+), 28 deletions(-)

diff --git a/.github/workflows/durabletask.yml b/.github/workflows/durabletask.yml
index 57fa5635..273f9e0a 100644
--- a/.github/workflows/durabletask.yml
+++ b/.github/workflows/durabletask.yml
@@ -14,6 +14,30 @@ permissions:
   contents: read
 
 jobs:
+  async-runtime-windows:
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12", "3.13", "3.14", "3.15"]
+    runs-on: windows-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install runtime test dependencies
+        run: |
+          python -m pip install --upgrade "pip>=25.1"
+          python -m pip install --group test -e .
+      - name: Check asyncio, replay, and dispatch cleanup
+        run: >-
+          python -m pytest
+          tests\durabletask\test_asyncio_orchestrators.py
+          tests\durabletask\test_coroutine_orchestrators.py
+          tests\durabletask\test_orchestration_generator_cleanup.py
+          tests\durabletask\test_orchestration_dispatch_recovery.py
+          -q
+
   lint-and-unit-tests:
     runs-on: ubuntu-latest
     steps:
diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py
index 695054c9..3c4edc79 100644
--- a/durabletask/internal/asyncio_loop.py
+++ b/durabletask/internal/asyncio_loop.py
@@ -9,6 +9,7 @@
 import contextvars
 import math
 import logging
+import sys
 from collections import deque
 from collections.abc import Callable, Coroutine, Generator
 from datetime import datetime, timedelta
@@ -16,6 +17,7 @@
 from typing import Any, TypeVar, cast
 
 from durabletask import task
+from durabletask.internal.exceptions import is_control_flow_exception
 
 T = TypeVar("T")
 
@@ -130,9 +132,11 @@ def __init__(
         self._observers: dict[int, _CompletionObserver] = {}
         self._completions: deque[Callable[[], None]] = deque()
         self._timers: dict[_DurableTimerHandle, _CompletionObserver] = {}
+        self._asyncgens: dict[int, Any] = {}
         self._closed = False
         self._closing = False
         self._fatal: BaseException | None = None
+        self._cleanup_failure: BaseException | None = None
         self._root: asyncio.Task[Any] | None = None
 
     def time(self) -> float:
@@ -237,6 +241,8 @@ def _complete_timer(self, handle: _DurableTimerHandle) -> None:
             self._ready.append(handle)
 
     async def wait_durable(self, durable: task.Task[T]) -> T:
+        if self._closing:
+            raise RuntimeError("Dispatch cleanup cannot await durable work.")
         if durable.is_complete:
             return durable.get_result()
         future = self.create_future()
@@ -283,7 +289,8 @@ def resume(self) -> None:
                 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() deliberately propagates control-flow BaseExceptions to the
+            # existing executor recovery boundary, not a success-shaped fallback.
             result = self._root.result()
             raise StopIteration(result)
         if not self._has_durable_work():
@@ -297,33 +304,57 @@ def _has_durable_work(self) -> bool:
     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
+        previous_hooks = sys.get_asyncgen_hooks()
         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)
+        sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter, finalizer=self._asyncgen_finalizer)
         try:
+            steps = 0
+            cleanup_limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
             while self._ready or self._cancel_ready:
                 self._ready.extend(handle for _, _, handle in sorted(self._cancel_ready))
                 self._cancel_ready.clear()
                 handle = self._ready.popleft()
                 if not handle.cancelled():
-                    handle._run()
+                    try:
+                        handle._run()
+                    except BaseException as error:
+                        if not self._closing:
+                            raise
+                        # Defer all exits until other cleanup tasks have drained.
+                        self._record_cleanup_error(error)
                 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)
+                    self._record_cleanup_error(error)
+                if not self._closing and self._is_complete():
+                    break
+                if not self._closing and self._root is not None and self._root.done():
+                    break
                 if not self._closing:
-                    if self._is_complete():
-                        break
-                    if self._root is not None and self._root.done():
-                        break
                     self._deliver_completions()
+                steps += 1
+                if self._closing and steps >= cleanup_limit:
+                    self._record_cleanup_error(RuntimeError(
+                        "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
+                    self._ready.clear()
+                    self._cancel_ready.clear()
+                    break
         finally:
+            sys.set_asyncgen_hooks(*previous_hooks)
             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 _asyncgen_firstiter(self, generator: Any) -> None:
+        self._asyncgens[id(generator)] = generator
+
+    def _asyncgen_finalizer(self, generator: Any) -> None:
+        if not self._closed:
+            self.create_task(generator.aclose())
+
     def call_exception_handler(self, context: dict[str, Any]) -> None:
         error = context.get("exception")
         if isinstance(error, StopIteration):
@@ -332,26 +363,47 @@ def call_exception_handler(self, context: dict[str, Any]) -> None:
             error = failure
         self._fatal = error if isinstance(error, BaseException) else RuntimeError(context["message"])
 
+    def _record_cleanup_error(self, error: BaseException, source: str = "cleanup callback") -> None:
+        self._logger.warning("%s: Async %s failed.", self.ctx.instance_id, source, exc_info=error)
+        if not isinstance(error, (Exception, asyncio.CancelledError)):
+            if self._cleanup_failure is None or is_control_flow_exception(error):
+                self._cleanup_failure = error
+
     def close(self) -> None:
-        """Cancel unfinished coroutines so their cleanup runs before the dispatch ends."""
-        if self._closed:
-            return
         self._closing = True
         try:
+            completed_at_close = {future for future in self._tasks if future.done()}
+            unobserved_at_close = {
+                future for future in completed_at_close if getattr(future, "_log_traceback", False)}
             for future in tuple(self._tasks):
                 if not future.done():
                     future.cancel()
             self._drain()
-            failure: BaseException | None = None
+            for generator in tuple(self._asyncgens.values()):
+                self.create_task(generator.aclose())
+            self._drain()
+            for future in tuple(self._tasks):
+                if not future.done():
+                    future.cancel()
+            self._drain()
+            failure: BaseException | None = self._cleanup_failure
             for future in tuple(self._tasks):
                 if not future.done():
-                    future.get_coro().close()
+                    try:
+                        future.get_coro().close()
+                    except BaseException as error:
+                        # Close siblings before propagating process-control exits.
+                        self._record_cleanup_error(error, "coroutine cleanup")
+                        failure = failure or error
                     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())
+                elif not future.cancelled():
+                    error = future.exception()
+                    if error is not None and (future not in completed_at_close or future in unobserved_at_close):
+                        self._record_cleanup_error(
+                            error, "background task" if future in completed_at_close else "task cleanup")
+            if self._cleanup_failure is not None:
+                raise self._cleanup_failure
             if failure is not None:
                 raise failure
         finally:
@@ -365,6 +417,7 @@ def close(self) -> None:
             self._waits.clear()
             self._timers.clear()
             self._tasks.clear()
+            self._asyncgens.clear()
 
     def run_in_executor(self, executor: Any, func: Any, *args: Any) -> asyncio.Future[Any]:
         raise RuntimeError("Threads/executors are unsupported in orchestrators; use an activity.")
diff --git a/durabletask/internal/exceptions.py b/durabletask/internal/exceptions.py
index 549a54f1..cfc787f4 100644
--- a/durabletask/internal/exceptions.py
+++ b/durabletask/internal/exceptions.py
@@ -1,9 +1,20 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+import asyncio
+from typing import cast
+
 import durabletask.internal.orchestrator_service_pb2 as pb
 
 
+def is_control_flow_exception(error: BaseException) -> bool:
+    control_flow_types = (asyncio.CancelledError, GeneratorExit, KeyboardInterrupt, SystemExit)
+    if isinstance(error, BaseExceptionGroup):
+        group = cast(BaseExceptionGroup[BaseException], error)
+        return group.subgroup(control_flow_types) is not None
+    return isinstance(error, control_flow_types)
+
+
 class VersionFailureException(Exception):
     def __init__(self, error_details: pb.TaskFailureDetails) -> None:
         super().__init__()
diff --git a/durabletask/worker.py b/durabletask/worker.py
index fa1c5f67..9dbb5ba6 100644
--- a/durabletask/worker.py
+++ b/durabletask/worker.py
@@ -63,14 +63,6 @@
 _WorkItem = tuple[Callable[..., Any], Callable[..., Any], tuple[Any, ...], dict[str, Any]]
 
 
-def _is_control_flow_exception(error: BaseException) -> bool:
-    control_flow_types = (asyncio.CancelledError, GeneratorExit, KeyboardInterrupt, SystemExit)
-    if isinstance(error, BaseExceptionGroup):
-        group = cast(BaseExceptionGroup[BaseException], error)
-        return group.subgroup(control_flow_types) is not None
-    return isinstance(error, control_flow_types)
-
-
 class ConcurrencyOptions:
     """Configuration options for controlling concurrency of different work item types and the thread pool size.
 
@@ -1671,6 +1663,7 @@ def __init__(self,
         self._exception_properties_provider = exception_properties_provider
         self._logger = logger
         self._async_loop: DurableEventLoop | None = None
+        self._async_cleanup = False
 
     def run(self, orchestrator: Generator[task.Task[Any], Any, Any] | Coroutine[Any, Any, Any]) -> None:
         if isinstance(orchestrator, Coroutine):
@@ -1683,6 +1676,7 @@ def run(self, orchestrator: Generator[task.Task[Any], Any, Any] | Coroutine[Any,
 
     def close(self) -> None:
         if self._async_loop is not None:
+            self._async_cleanup = True
             self._async_loop.close()
         elif self._generator is not None:
             self._generator.close()
@@ -1849,6 +1843,8 @@ def get_actions(self) -> list[pb.OrchestratorAction]:
         return current_actions
 
     def next_sequence_number(self) -> int:
+        if self._async_cleanup:
+            raise RuntimeError("Dispatch cleanup cannot schedule durable work; use normal orchestrator control flow.")
         self._sequence_number += 1
         return self._sequence_number
 
@@ -2493,7 +2489,7 @@ def execute(
                 ctx.set_failed(ex)
 
             except BaseException as ex:
-                if _is_control_flow_exception(ex):
+                if pe.is_control_flow_exception(ex):
                     raise
                 self._logger.error(
                     f"{instance_id}: Orchestrator {orchestration_name} exited unexpectedly.",
@@ -2527,7 +2523,7 @@ def execute(
             result = ExecutionResults(
                 actions=actions, encoded_custom_status=ctx._encoded_custom_status,  # pyright: ignore[reportPrivateUsage]
                 orchestration_trace_context=ctx._orchestration_trace_context,  # pyright: ignore[reportPrivateUsage]
-                # Snapshot before generator teardown can exit a critical section.
+                # Snapshot before orchestrator teardown can exit a critical section.
                 entity_lock_releases=list(ctx._entity_context.get_lock_release_messages()),  # pyright: ignore[reportPrivateUsage]
             )
             if unexpected_failure and result.entity_lock_releases:
@@ -2560,7 +2556,7 @@ def execute(
                     exc_info=True,
                 )
                 if not dispatch_exited:
-                    if _is_control_flow_exception(ex) or result is None:
+                    if pe.is_control_flow_exception(ex) or result is None:
                         raise
                     # Keep an existing failure rather than mask its cause.
                     if not any(
diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py
index d0c2db8f..18cd0036 100644
--- a/tests/durabletask/test_asyncio_orchestrators.py
+++ b/tests/durabletask/test_asyncio_orchestrators.py
@@ -9,6 +9,7 @@
 import os
 import subprocess
 import sys
+from contextlib import asynccontextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
 from datetime import timedelta
@@ -335,6 +336,49 @@ async def branch(index: int) -> None:
     assert json.loads(_completion(_execute(orchestrator, old, new)).result.value) == "done"
 
 
+@pytest.mark.parametrize("under_loop", [False, True])
+def test_dispatch_cleanup_unwinds_scopes_contexts_asyncgens(under_loop: bool) -> None:
+    cleaned: list[str] = []
+    references: list[Any] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        references.extend([ref(ctx), ref(asyncio.get_running_loop()), ref(asyncio.current_task())])
+
+        async def resource():
+            try:
+                yield 1
+            finally:
+                await asyncio.sleep(0)
+                cleaned.append("asyncgen")
+
+        opened = resource()
+        await anext(opened)
+        try:
+            async with asyncio.TaskGroup():
+                await ctx.call_activity("wait")
+        finally:
+            await asyncio.sleep(0)
+            cleaned.append("root")
+
+    def dispatch() -> None:
+        result = _execute(orchestrator, [], _start())
+        assert [a.scheduleTask.name for a in result.actions] == ["wait"]
+        assert cleaned == ["root", "asyncgen"]
+
+    async def outer() -> None:
+        before = asyncio.current_task()
+        dispatch()
+        assert asyncio.current_task() is before
+        await asyncio.sleep(0)
+
+    if under_loop:
+        asyncio.run(outer())
+    else:
+        dispatch()
+    gc.collect()
+    assert all(reference() is None for reference in references)
+
+
 def test_shared_operation_survives_local_waiter_cancellation() -> None:
     async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
         durable = ctx.call_activity("shared", return_type=str)
@@ -375,6 +419,57 @@ async def waiter() -> None:
     assert json.loads(_completion(result).result.value) == "cancelled"
 
 
+@pytest.mark.parametrize("terminal", ["suspend", "terminate", "continue", "history-failure"])
+def test_async_teardown_preserves_response_and_cannot_schedule_work(
+        terminal: str, caplog: pytest.LogCaptureFixture) -> None:
+    cleaned: list[str] = []
+
+    @asynccontextmanager
+    async def resource():
+        try:
+            yield
+        finally:
+            await asyncio.sleep(0)
+            cleaned.append("resource")
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async with resource():
+            try:
+                ctx.set_custom_status("waiting")
+                if terminal == "continue":
+                    ctx.continue_as_new("next")
+                await ctx.call_activity("wait")
+            finally:
+                try:
+                    await asyncio.sleep(0)
+                    cleaned.append("inner")
+                    await ctx.call_activity("forbidden-cleanup")
+                finally:
+                    await asyncio.sleep(0)
+                    cleaned.append("outer")
+
+    old = _start()
+    if terminal == "terminate":
+        new = [pb.HistoryEvent(executionTerminated=pb.ExecutionTerminatedEvent())]
+    elif terminal == "history-failure":
+        new = [helpers.new_task_scheduled_event(1, "wrong")]
+    else:
+        old, new = [], old
+    result = _execute(orchestrator, old, new)
+    assert cleaned == ["inner", "outer", "resource"]
+    assert result.encoded_custom_status == '"waiting"'
+    assert all(action.scheduleTask.name != "forbidden-cleanup" for action in result.actions)
+    if terminal == "terminate":
+        assert _completion(result).orchestrationStatus == pb.ORCHESTRATION_STATUS_TERMINATED
+    elif terminal == "history-failure":
+        assert _completion(result).failureDetails.errorType.endswith("NonDeterminismError")
+    elif terminal == "continue":
+        assert _completion(result).orchestrationStatus == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW
+    else:
+        assert [action.scheduleTask.name for action in result.actions] == ["wait"]
+    assert "Dispatch cleanup cannot schedule durable work" in caplog.text
+
+
 def test_unhandled_local_root_cancellation_fails_not_abandons() -> None:
     async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
         raise asyncio.CancelledError
@@ -498,6 +593,42 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
         "use an activity", "ordinary activity", "different loop"))
 
 
+@pytest.mark.parametrize("async_generator", [False, True])
+@pytest.mark.parametrize("unexpected", [False, True])
+def test_cleanup_failure_policy_includes_async_generator_tasks(
+        async_generator: bool, unexpected: bool, caplog: pytest.LogCaptureFixture) -> None:
+    class CleanupAbort(BaseException):
+        pass
+
+    error = CleanupAbort("cleanup failed") if unexpected else ValueError("cleanup failed")
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def resource():
+            try:
+                yield 1
+            finally:
+                await asyncio.sleep(0)
+                raise error
+
+        opened = resource()
+        await anext(opened)
+        try:
+            await ctx.call_activity("wait")
+        finally:
+            if not async_generator:
+                await opened.aclose()
+
+    result = _execute(orchestrator, [], _start())
+    if unexpected:
+        complete = _completion(result)
+        assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+        assert complete.failureDetails.errorMessage == "cleanup failed"
+        assert complete.failureDetails.errorType.endswith("CleanupAbort")
+    else:
+        assert [action.scheduleTask.name for action in result.actions] == ["wait"]
+        assert "cleanup failed" in caplog.text
+
+
 def test_task_group_cancellation_durable_cleanup_has_stable_action_order() -> None:
     async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
         async def branch(index: int) -> None:
@@ -672,6 +803,80 @@ async def fail() -> None:
     assert ("background failure" in caplog.text) is not observed
 
 
+def test_unexpected_cleanup_failure_does_not_hide_other_task_errors(
+        caplog: pytest.LogCaptureFixture) -> None:
+    class CleanupAbort(BaseException):
+        pass
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def branch(error: BaseException) -> None:
+            try:
+                await ctx.wait_for_external_event("never")
+            finally:
+                await asyncio.sleep(0)
+                raise error
+
+        asyncio.create_task(branch(CleanupAbort("first cleanup failure")))
+        asyncio.create_task(branch(ValueError("second cleanup failure")))
+        await ctx.call_activity("wait")
+
+    complete = _completion(_execute(orchestrator, [], _start()))
+    assert complete.failureDetails.errorMessage == "first cleanup failure"
+    assert "second cleanup failure" in caplog.text
+
+
+@pytest.mark.parametrize("control_flow", [False, True])
+def test_cleanup_callback_failure_still_drains_tasks_and_asyncgens(
+        control_flow: bool, caplog: pytest.LogCaptureFixture) -> None:
+    cleaned: list[str] = []
+    retained: list[asyncio.Task[Any]] = []
+
+    def fail() -> None:
+        raise SystemExit("cleanup exit") if control_flow else ValueError("cleanup callback")
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def resource():
+            try:
+                yield 1
+            finally:
+                await asyncio.sleep(0)
+                cleaned.append("asyncgen")
+
+        async def sibling() -> None:
+            try:
+                await ctx.wait_for_external_event("never")
+            finally:
+                await asyncio.sleep(0)
+                cleaned.append("sibling")
+
+        opened = resource()
+        await anext(opened)
+        retained.append(asyncio.create_task(sibling()))
+        current = asyncio.current_task()
+        assert current is not None
+        retained.append(current)
+        try:
+            await ctx.call_activity("wait")
+        finally:
+            if control_flow:
+                fail()
+            asyncio.get_running_loop().call_soon(fail)
+            await asyncio.sleep(0)
+            cleaned.append("root")
+
+    if control_flow:
+        with pytest.raises(SystemExit, match="cleanup exit"):
+            _execute(orchestrator, [], _start())
+    else:
+        result = _execute(orchestrator, [], _start())
+        assert [action.scheduleTask.name for action in result.actions] == ["wait"]
+        assert "cleanup callback" in caplog.text
+        assert "root" in cleaned
+    assert "sibling" in cleaned
+    assert "asyncgen" in cleaned
+    assert all(pending.done() for pending in retained)
+
+
 @pytest.mark.parametrize("count", [128, 4096])
 def test_fanout_waiter_inspections_scale_linearly(
         count: int, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None:
@@ -696,6 +901,72 @@ async def branch(index: int) -> Any:
     assert not caplog.records
 
 
+def test_nonquiescent_cleanup_does_not_skip_asyncgen_closure(
+        caplog: pytest.LogCaptureFixture) -> None:
+    cleaned: list[str] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def resource():
+            try:
+                yield 1
+            finally:
+                await asyncio.sleep(0)
+                cleaned.append("asyncgen")
+
+        opened = resource()
+        await anext(opened)
+        try:
+            await ctx.call_activity("wait")
+        finally:
+            while True:
+                await asyncio.sleep(0)
+
+    result = _execute(orchestrator, [], _start())
+    assert [action.scheduleTask.name for action in result.actions] == ["wait"]
+    assert cleaned == ["asyncgen"]
+    assert "cleanup did not quiesce" in caplog.text
+
+
+@pytest.mark.parametrize("grouped", [False, True])
+@pytest.mark.parametrize("custom_error", [False, True])
+def test_forced_cleanup_preserves_process_control_after_other_errors(grouped: bool, custom_error: bool) -> None:
+    class CleanupAbort(BaseException):
+        pass
+
+    cleaned: list[str] = []
+    exit_error: BaseException = SystemExit("forced cleanup exit")
+    if grouped:
+        exit_error = BaseExceptionGroup("forced cleanup group", [exit_error, ValueError("nested")])
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def branch(name: str, error: BaseException | None) -> None:
+            try:
+                await ctx.wait_for_external_event(name)
+            finally:
+                try:
+                    while True:
+                        try:
+                            await asyncio.Future()
+                        except asyncio.CancelledError:
+                            # Force teardown to close a cancellation-resistant coroutine.
+                            pass
+                finally:
+                    cleaned.append(name)
+                    if error is not None:
+                        raise error
+
+        first_error = CleanupAbort("custom cleanup error") if custom_error else ValueError("ordinary cleanup error")
+        asyncio.create_task(branch("ordinary", first_error))
+        asyncio.create_task(branch("exit", exit_error))
+        asyncio.create_task(branch("last", None))
+        await ctx.call_activity("wait")
+
+    with pytest.raises(type(exit_error)) as caught:
+        _execute(orchestrator, [], _start())
+    assert caught.value is exit_error
+    assert cleaned == ["ordinary", "exit", "last"]
+
+
 @pytest.mark.parametrize("timer", [False, True])
 def test_callback_stop_iteration_fails_instead_of_completing_orchestrator(timer: bool) -> None:
     async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
@@ -758,3 +1029,31 @@ def callback() -> None:
         else:
             assert json.loads(_completion(result).result.value) == len(remaining)
         old += new
+
+
+def test_cooperative_cleanup_prioritizes_grouped_process_control() -> None:
+    class CleanupAbort(BaseException):
+        pass
+
+    exit_error = BaseExceptionGroup("cleanup exit", [SystemExit("stop"), ValueError("nested")])
+    cleaned: list[str] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def branch(name: str, error: BaseException | None) -> None:
+            try:
+                await ctx.wait_for_external_event(name)
+            finally:
+                await asyncio.sleep(0)
+                cleaned.append(name)
+                if error is not None:
+                    raise error
+
+        asyncio.create_task(branch("custom", CleanupAbort("first failure")))
+        asyncio.create_task(branch("exit", exit_error))
+        asyncio.create_task(branch("last", None))
+        await ctx.call_activity("wait")
+
+    with pytest.raises(BaseExceptionGroup) as caught:
+        _execute(orchestrator, [], _start())
+    assert caught.value is exit_error
+    assert cleaned == ["custom", "exit", "last"]

From de57073201ac24e8a80edfccdcbe56a9bb280480 Mon Sep 17 00:00:00 2001
From: Andy Staples <andystaples@microsoft.com>
Date: Thu, 12 Nov 2026 11:12:18 -0700
Subject: [PATCH 2/4] Harden async dispatch cleanup ownership and signal
 precedence

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 docs/async-orchestrators.md                   |  11 +-
 durabletask/internal/asyncio_loop.py          |  43 +++++-
 durabletask/worker.py                         |   4 +
 .../durabletask/test_asyncio_orchestrators.py | 132 +++++++++++++++++-
 4 files changed, 183 insertions(+), 7 deletions(-)

diff --git a/docs/async-orchestrators.md b/docs/async-orchestrators.md
index 50cbf8a5..273105b4 100644
--- a/docs/async-orchestrators.md
+++ b/docs/async-orchestrators.md
@@ -250,13 +250,20 @@ Unexpected non-cancellation `BaseException`
 exits retain the existing terminal-failure/recovery policy; an existing primary
 failure is preserved. Teardown status/actions do not replace the response.
 Held entity-release obligations are captured before teardown for terminal
-payload/failure handling.
+payload/failure handling. Exiting an entity critical section during teardown
+does not schedule additional release messages, and a `SystemExit` or
+`KeyboardInterrupt` raised by any finalizer takes precedence over other
+cleanup failures.
 
 Keep workflow cleanup that must be durable in **normal orchestration control
 flow**, before return or continue-as-new, not solely in a dispatch finalizer.
 No cleanup is guaranteed after process crash or forced termination.
 Async-generator *helpers* are closed; async-generator *orchestrators*
-(`async def` containing `yield`) are rejected.
+(`async def` containing `yield`) are rejected. The dispatch owns every async
+generator it starts until teardown, including generators first iterated by
+another generator's cleanup. To run a generator's cleanup earlier, close it
+explicitly with `await gen.aclose()` or iterate it inside
+`contextlib.aclosing(...)`.
 
 ## Replay and framework adapters
 
diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py
index 3c4edc79..341fc737 100644
--- a/durabletask/internal/asyncio_loop.py
+++ b/durabletask/internal/asyncio_loop.py
@@ -25,6 +25,13 @@
 _MICROSECOND = timedelta(microseconds=1)
 
 
+def _is_process_control(error: BaseException | None) -> bool:
+    if isinstance(error, BaseExceptionGroup):
+        group = cast(BaseExceptionGroup[BaseException], error)
+        return group.subgroup((KeyboardInterrupt, SystemExit)) is not None
+    return isinstance(error, (KeyboardInterrupt, SystemExit))
+
+
 class _CompletionObserver(task.CompositeTask[None]):
     def __init__(self, durable: task.Task[Any], callback: Callable[[], None]) -> None:
         self._callback = callback
@@ -133,6 +140,7 @@ def __init__(
         self._completions: deque[Callable[[], None]] = deque()
         self._timers: dict[_DurableTimerHandle, _CompletionObserver] = {}
         self._asyncgens: dict[int, Any] = {}
+        self._asyncgen_prune_at = 64
         self._closed = False
         self._closing = False
         self._fatal: BaseException | None = None
@@ -301,7 +309,7 @@ def resume(self) -> None:
     def _has_durable_work(self) -> bool:
         return any(not durable.is_complete for durable, _ in self._waits.values()) or bool(self._timers)
 
-    def _drain(self) -> None:
+    def _drain(self) -> int:
         previous_loop = asyncio._get_running_loop()
         previous_task = asyncio.current_task(previous_loop) if previous_loop is not None else None
         previous_hooks = sys.get_asyncgen_hooks()
@@ -342,6 +350,7 @@ def _drain(self) -> None:
                     self._ready.clear()
                     self._cancel_ready.clear()
                     break
+            return steps
         finally:
             sys.set_asyncgen_hooks(*previous_hooks)
             asyncio._set_running_loop(previous_loop)
@@ -349,6 +358,11 @@ def _drain(self) -> None:
                 asyncio.tasks._enter_task(previous_loop, previous_task)
 
     def _asyncgen_firstiter(self, generator: Any) -> None:
+        if len(self._asyncgens) >= self._asyncgen_prune_at:
+            # Exhausted generators have no cleanup left; amortized pruning bounds retention.
+            self._asyncgens = {
+                key: value for key, value in self._asyncgens.items() if value.ag_frame is not None}
+            self._asyncgen_prune_at = max(64, 2 * len(self._asyncgens))
         self._asyncgens[id(generator)] = generator
 
     def _asyncgen_finalizer(self, generator: Any) -> None:
@@ -366,9 +380,30 @@ def call_exception_handler(self, context: dict[str, Any]) -> None:
     def _record_cleanup_error(self, error: BaseException, source: str = "cleanup callback") -> None:
         self._logger.warning("%s: Async %s failed.", self.ctx.instance_id, source, exc_info=error)
         if not isinstance(error, (Exception, asyncio.CancelledError)):
+            if _is_process_control(self._cleanup_failure):
+                return
             if self._cleanup_failure is None or is_control_flow_exception(error):
                 self._cleanup_failure = error
 
+    def _close_asyncgens(self, closed: dict[int, Any]) -> None:
+        # Closing one generator may first-iterate another; sweep until no new generators appear.
+        # Closed generators stay referenced so their ids cannot be reused during cleanup.
+        steps = 0
+        while True:
+            pending = [
+                generator for key, generator in self._asyncgens.items() if closed.get(key) is not generator]
+            if not pending:
+                return
+            limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
+            if steps >= limit:
+                self._record_cleanup_error(RuntimeError(
+                    "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
+                return
+            for generator in pending:
+                closed[id(generator)] = generator
+                self.create_task(generator.aclose())
+            steps += max(1, self._drain())
+
     def close(self) -> None:
         self._closing = True
         try:
@@ -379,13 +414,13 @@ def close(self) -> None:
                 if not future.done():
                     future.cancel()
             self._drain()
-            for generator in tuple(self._asyncgens.values()):
-                self.create_task(generator.aclose())
-            self._drain()
+            closed_asyncgens: dict[int, Any] = {}
+            self._close_asyncgens(closed_asyncgens)
             for future in tuple(self._tasks):
                 if not future.done():
                     future.cancel()
             self._drain()
+            self._close_asyncgens(closed_asyncgens)
             failure: BaseException | None = self._cleanup_failure
             for future in tuple(self._tasks):
                 if not future.done():
diff --git a/durabletask/worker.py b/durabletask/worker.py
index 9dbb5ba6..df31c42c 100644
--- a/durabletask/worker.py
+++ b/durabletask/worker.py
@@ -2234,6 +2234,10 @@ def _exit_critical_section(self) -> None:
             # _exit_critical_section both from the EntityLock and the continue_as_new logic. We must keep both calls in
             # case the user code crashes after calling continue_as_new but before the EntityLock object is exited.
             return
+        if self._async_cleanup:
+            # Lock disposal during local teardown is not a durable release; the
+            # pre-teardown response already carries any release this dispatch owns.
+            return
         for entity_unlock_message in self._entity_context.emit_lock_release_messages():
             task_id = self.next_sequence_number()
             action = pb.OrchestratorAction(id=task_id, sendEntityMessage=entity_unlock_message)
diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py
index 18cd0036..1c9f25d7 100644
--- a/tests/durabletask/test_asyncio_orchestrators.py
+++ b/tests/durabletask/test_asyncio_orchestrators.py
@@ -9,6 +9,7 @@
 import os
 import subprocess
 import sys
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -20,7 +21,7 @@
 
 import pytest
 
-from durabletask import task
+from durabletask import entities, task
 from durabletask.internal import helpers
 from durabletask.internal import orchestrator_service_pb2 as pb
 from tests.durabletask.test_coroutine_orchestrators import (
@@ -1057,3 +1058,132 @@ async def branch(name: str, error: BaseException | None) -> None:
         _execute(orchestrator, [], _start())
     assert caught.value is exit_error
     assert cleaned == ["custom", "exit", "last"]
+
+
+def _cleanup_signal(kind: str) -> BaseException:
+    if kind == "system_exit":
+        return SystemExit("stop")
+    if kind == "keyboard_interrupt":
+        return KeyboardInterrupt("interrupt")
+    if kind == "nested_group":
+        return BaseExceptionGroup("outer", [BaseExceptionGroup("inner", [SystemExit("stop")])])
+    if kind == "generator_exit":
+        return GeneratorExit("closed")
+    return BaseExceptionGroup("cancelled", [asyncio.CancelledError()])
+
+
+@pytest.mark.parametrize("control_first", [True, False])
+@pytest.mark.parametrize("other", ["generator_exit", "cancelled_group"])
+@pytest.mark.parametrize("control", ["system_exit", "keyboard_interrupt", "nested_group"])
+def test_cleanup_keeps_process_control_over_other_control_flow(
+        control: str, other: str, control_first: bool) -> None:
+    control_error = _cleanup_signal(control)
+    errors = [control_error, _cleanup_signal(other)]
+    if not control_first:
+        errors.reverse()
+    cleaned: list[int] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def branch(index: int) -> None:
+            try:
+                await ctx.wait_for_external_event(f"never-{index}")
+            finally:
+                cleaned.append(index)
+                raise errors[index]
+
+        asyncio.create_task(branch(0))
+        asyncio.create_task(branch(1))
+        await ctx.call_activity("wait")
+
+    with pytest.raises(type(control_error)) as caught:
+        _execute(orchestrator, [], _start())
+    assert caught.value is control_error
+    assert cleaned == [0, 1]
+
+
+def test_cleanup_closes_async_generators_started_during_generator_cleanup(
+        caplog: pytest.LogCaptureFixture) -> None:
+    closed: list[str] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def child() -> AsyncIterator[int]:
+            try:
+                yield 1
+            finally:
+                await asyncio.sleep(0)
+                closed.append("child")
+
+        async def outer() -> AsyncIterator[int]:
+            try:
+                yield 1
+            finally:
+                await anext(child())
+                closed.append("outer")
+
+        await anext(outer())
+        await ctx.call_activity("wait")
+
+    actions = _execute(orchestrator, [], _start()).actions
+    assert [action.scheduleTask.name for action in actions] == ["wait"]
+    assert closed == ["outer", "child"]
+    assert "Async" not in caplog.text
+
+
+def test_exhausted_async_generators_are_not_retained_until_dispatch_end() -> None:
+    references: list[ReferenceType[Any]] = []
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def values() -> AsyncIterator[int]:
+            yield 1
+
+        for _ in range(200):
+            generator = values()
+            references.append(ref(generator))
+            async for _ in generator:
+                pass
+            del generator
+        gc.collect()
+        assert sum(reference() is not None for reference in references) < 100
+        await ctx.call_activity("wait")
+
+    _execute(orchestrator, [], _start())
+
+
+@pytest.mark.parametrize("suspend", ["activity", "entity"])
+def test_entity_lock_disposal_during_teardown_does_not_schedule_releases(
+        suspend: str, caplog: pytest.LogCaptureFixture) -> None:
+    first = entities.EntityInstanceId("counter", "a")
+    second = entities.EntityInstanceId("counter", "b")
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any:
+        with await ctx.lock_entities([first, second]):
+            if suspend == "activity":
+                return await ctx.call_activity("inside-lock")
+            return await ctx.call_entity(first, "get")
+
+    old = _start()
+    lock = _execute(orchestrator, [], old).actions[0].sendEntityMessage.entityLockRequested
+    old += [helpers.new_event_sent_event(1, str(first), json.dumps({"id": lock.criticalSectionId}))]
+    new = [helpers.new_event_raised_event(lock.criticalSectionId, None)]
+    actions = _execute(orchestrator, old, new).actions
+    assert len(actions) == 1
+    if suspend == "activity":
+        assert actions[0].scheduleTask.name == "inside-lock"
+        old += new + [helpers.new_task_scheduled_event(2, "inside-lock")]
+        new = [helpers.new_task_completed_event(2, json.dumps(42))]
+    else:
+        call = actions[0].sendEntityMessage.entityOperationCalled
+        assert call.targetInstanceId.value == str(first)
+        old += new + [helpers.new_event_sent_event(2, str(first), json.dumps({"id": call.requestId}))]
+        new = [helpers.new_event_raised_event(call.requestId, json.dumps({"result": "42"}))]
+    assert "cleanup" not in caplog.text
+
+    result = _execute(orchestrator, old, new)
+    actions = result.actions
+    releases = [action.sendEntityMessage.entityUnlockSent for action in actions[:-1]]
+    # Entity calls inside a critical section keep the legacy shared-list release order.
+    expected = [first, second] if suspend == "activity" else [second, first]
+    assert [release.targetInstanceId.value for release in releases] == [str(e) for e in expected]
+    assert [action.id for action in actions] == [3, 4, 5]
+    assert json.loads(_completion(result).result.value) == 42
+    assert "cleanup" not in caplog.text

From c164cb697613b02067c17c77b7e97022a73a09b8 Mon Sep 17 00:00:00 2001
From: Andy Staples <andystaples@microsoft.com>
Date: Thu, 12 Nov 2026 11:33:56 -0700
Subject: [PATCH 3/4] Bound async-generator cleanup discovery budget

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 CHANGELOG.md                                  |  3 +
 docs/async-orchestrators.md                   |  5 +-
 durabletask-azuremanaged/CHANGELOG.md         |  3 +
 durabletask/internal/asyncio_loop.py          |  3 +-
 .../durabletask/test_asyncio_orchestrators.py | 92 ++++++++++++++++---
 5 files changed, 89 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3604c296..62e756b1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,9 @@ FIXED
 when an intermediate async wrapper decorates a synchronous function, including
 callable instances and partials. Cyclic wrapper chains fail instead of hanging
 worker dispatch.
+- Async dispatch cleanup no longer hangs when async-generator finalizers
+continually start successor generators; exceeding the cleanup budget is reported
+without changing the snapshotted response.
 - Coroutine task-group cancellation now schedules shared-Future callbacks
 deterministically across replays, including when a dispatch resumes on another
 worker using the same SDK and Python runtime. Previously recorded cancellation
diff --git a/docs/async-orchestrators.md b/docs/async-orchestrators.md
index 273105b4..ac386350 100644
--- a/docs/async-orchestrators.md
+++ b/docs/async-orchestrators.md
@@ -242,7 +242,10 @@ dispatch boundary**, before workflow completion. Local cleanup can await ready
 work such as `sleep(0)`; it must not wait for new durable or external work.
 Scheduling/awaiting durable work during async teardown raises an actionable
 error. Teardown must cooperate with cancellation; a non-quiescent cleanup
-queue is bounded and reported, not allowed to run indefinitely.
+queue is bounded and reported, not allowed to run indefinitely. Each
+async-generator sweep has a step budget fixed at sweep entry; tasks or generators
+created during that sweep cannot increase it. This bounds cooperative scheduling
+steps, not wall-clock time or code that never yields.
 
 Ordinary teardown failures (including `CancelledError`) are logged without
 changing the decided response or skipping other cooperative finalizers.
diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md
index 11b59946..157a4604 100644
--- a/durabletask-azuremanaged/CHANGELOG.md
+++ b/durabletask-azuremanaged/CHANGELOG.md
@@ -9,6 +9,9 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
+- With the corresponding core SDK update, async dispatch cleanup no longer
+hangs when async-generator finalizers continually start successor generators.
+Exceeding the cleanup budget is reported without changing the snapshotted response.
 - With the corresponding core SDK update, coroutine task-group cancellation
 schedules shared-Future callbacks deterministically when replayed on another
 worker using the same SDK and Python runtime. See the
diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py
index 341fc737..2136d970 100644
--- a/durabletask/internal/asyncio_loop.py
+++ b/durabletask/internal/asyncio_loop.py
@@ -389,12 +389,13 @@ def _close_asyncgens(self, closed: dict[int, Any]) -> None:
         # Closing one generator may first-iterate another; sweep until no new generators appear.
         # Closed generators stay referenced so their ids cannot be reused during cleanup.
         steps = 0
+        # Cleanup-created tasks and generators must not replenish the sweep's budget.
+        limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
         while True:
             pending = [
                 generator for key, generator in self._asyncgens.items() if closed.get(key) is not generator]
             if not pending:
                 return
-            limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
             if steps >= limit:
                 self._record_cleanup_error(RuntimeError(
                     "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py
index 1c9f25d7..3fcf1b9b 100644
--- a/tests/durabletask/test_asyncio_orchestrators.py
+++ b/tests/durabletask/test_asyncio_orchestrators.py
@@ -1101,32 +1101,91 @@ async def branch(index: int) -> None:
     assert cleaned == [0, 1]
 
 
+@pytest.mark.parametrize("count", [2, 32, 6000])
 def test_cleanup_closes_async_generators_started_during_generator_cleanup(
-        caplog: pytest.LogCaptureFixture) -> None:
-    closed: list[str] = []
+        count: int, caplog: pytest.LogCaptureFixture) -> None:
+    closed: list[int] = []
 
     async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
-        async def child() -> AsyncIterator[int]:
+        async def resource(index: int) -> AsyncIterator[int]:
             try:
                 yield 1
             finally:
+                if index + 1 < count:
+                    await anext(resource(index + 1))
                 await asyncio.sleep(0)
-                closed.append("child")
-
-        async def outer() -> AsyncIterator[int]:
-            try:
-                yield 1
-            finally:
-                await anext(child())
-                closed.append("outer")
+                closed.append(index)
 
-        await anext(outer())
+        await anext(resource(0))
         await ctx.call_activity("wait")
 
     actions = _execute(orchestrator, [], _start()).actions
     assert [action.scheduleTask.name for action in actions] == ["wait"]
-    assert closed == ["outer", "child"]
-    assert "Async" not in caplog.text
+    assert closed == list(range(count))
+    if count == 6000:
+        assert "Async cleanup did not quiesce" in caplog.text
+    else:
+        assert "Async" not in caplog.text
+
+
+@pytest.mark.parametrize("control", [None, "system_exit", "keyboard_interrupt", "nested_group"])
+def test_infinite_asyncgen_discovery_has_bounded_cleanup(control: str | None) -> None:
+    code = """
+import asyncio
+import logging
+import sys
+from durabletask import task
+from tests.durabletask.test_asyncio_orchestrators import _cleanup_signal
+from tests.durabletask.test_coroutine_orchestrators import _execute, _start
+
+control = sys.argv[1]
+exit_error = None if control == "none" else _cleanup_signal(control)
+warnings = []
+
+class Capture(logging.Handler):
+    def emit(self, record):
+        warnings.append(logging.Formatter().format(record))
+
+logging.getLogger().addHandler(Capture())
+
+async def orchestrator(ctx: task.OrchestrationContext, _):
+    async def resource():
+        try:
+            yield 1
+        finally:
+            await anext(resource())
+
+    async def signal_on_forced_close():
+        try:
+            await ctx.wait_for_external_event("never")
+        finally:
+            try:
+                while True:
+                    try:
+                        await asyncio.Future()
+                    except asyncio.CancelledError:
+                        pass
+            finally:
+                raise exit_error
+
+    if exit_error is not None:
+        asyncio.create_task(signal_on_forced_close())
+    await anext(resource())
+    await ctx.call_activity("wait")
+
+try:
+    result = _execute(orchestrator, [], _start())
+except BaseException as error:
+    assert exit_error is not None and error is exit_error, repr(error)
+else:
+    assert exit_error is None
+    assert [action.scheduleTask.name for action in result.actions] == ["wait"]
+assert any("Async cleanup did not quiesce" in message for message in warnings), warnings
+"""
+    process = subprocess.run(
+        [sys.executable, "-c", code, control or "none"],
+        capture_output=True, text=True, timeout=30)
+    assert process.returncode == 0, process.stdout + process.stderr
 
 
 def test_exhausted_async_generators_are_not_retained_until_dispatch_end() -> None:
@@ -1146,7 +1205,10 @@ async def values() -> AsyncIterator[int]:
         assert sum(reference() is not None for reference in references) < 100
         await ctx.call_activity("wait")
 
-    _execute(orchestrator, [], _start())
+    actions = _execute(orchestrator, [], _start()).actions
+    assert len(actions) == 1
+    assert actions[0].HasField("scheduleTask")
+    assert actions[0].scheduleTask.name == "wait"
 
 
 @pytest.mark.parametrize("suspend", ["activity", "entity"])

From 344196a643b1d2b0415c34ac3546f7b0b8ecde95 Mon Sep 17 00:00:00 2001
From: Andy Staples <andystaples@microsoft.com>
Date: Thu, 12 Nov 2026 16:07:52 -0700
Subject: [PATCH 4/4] Make async cleanup budgets cumulative and close
 idempotent

Spend each async-generator sweep's drains from the budget fixed at sweep
entry, report truncation once, keep the first custom cleanup failure unless a
process-control exit supersedes it, and restore the closed-loop guard. Add
regression tests for these and for restoring host async-generator hooks, and
document the cleanup cap prominently. Remove changelog entries for unreleased
behavior.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 CHANGELOG.md                                  |  3 -
 docs/async-orchestrators.md                   | 16 +++--
 durabletask-azuremanaged/CHANGELOG.md         |  3 -
 durabletask/internal/asyncio_loop.py          | 37 +++++++----
 .../durabletask/test_asyncio_orchestrators.py | 61 ++++++++++++++++++-
 5 files changed, 95 insertions(+), 25 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 55d4cc08..40aa2e88 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,9 +9,6 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
-- Async dispatch cleanup no longer hangs when async-generator finalizers
-continually start successor generators; exceeding the cleanup budget is reported
-without changing the snapshotted response.
 - 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/docs/async-orchestrators.md b/docs/async-orchestrators.md
index bddef2c6..5e362e38 100644
--- a/docs/async-orchestrators.md
+++ b/docs/async-orchestrators.md
@@ -302,11 +302,17 @@ continue as new, encounter a history failure, or reject a version.
 Local cleanup can await ready work such as `sleep(0)`; it must not wait for new
 durable or external work.
 Scheduling/awaiting durable work during async teardown raises an actionable
-error. Teardown must cooperate with cancellation; a non-quiescent cleanup
-queue is bounded and reported, not allowed to run indefinitely. Each
-async-generator sweep has a step budget fixed at sweep entry; tasks or generators
-created during that sweep cannot increase it. This bounds cooperative scheduling
-steps, not wall-clock time or code that never yields.
+error. Teardown must cooperate with cancellation.
+
+> [!WARNING]
+> Async teardown is capped at about 10,000 event-loop steps per cleanup phase
+> (more for dispatches with many tasks or async generators). Cleanup that is
+> still running when the cap is reached is truncated: remaining callbacks are
+> dropped, unfinished coroutines are closed, and a warning is logged. The
+> response decided before teardown is unchanged. Each async-generator sweep's
+> budget is fixed at sweep entry; tasks or generators created during that sweep
+> cannot increase it. The cap bounds cooperative scheduling steps, not
+> wall-clock time or code that never yields.
 
 Ordinary teardown failures (including `CancelledError`) are logged without
 changing the decided response or skipping other cooperative finalizers.
diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md
index d117aa06..1b327beb 100644
--- a/durabletask-azuremanaged/CHANGELOG.md
+++ b/durabletask-azuremanaged/CHANGELOG.md
@@ -9,9 +9,6 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 FIXED
 
-- With the corresponding core SDK update, async dispatch cleanup no longer
-hangs when async-generator finalizers continually start successor generators.
-Exceeding the cleanup budget is reported without changing the snapshotted response.
 - With the corresponding core SDK update, activity and sub-orchestration
 `retry_timeout` decisions use replay-safe orchestration time rather than
 wall-clock time. See the [core upgrade guidance](../CHANGELOG.md#unreleased)
diff --git a/durabletask/internal/asyncio_loop.py b/durabletask/internal/asyncio_loop.py
index 6a9fe76e..68054bff 100644
--- a/durabletask/internal/asyncio_loop.py
+++ b/durabletask/internal/asyncio_loop.py
@@ -39,7 +39,6 @@
 from typing import Any, TypeVar, cast
 
 from durabletask import task
-from durabletask.internal.exceptions import is_control_flow_exception
 
 T = TypeVar("T")
 
@@ -169,6 +168,7 @@ def __init__(
         self._closing = False
         self._fatal: BaseException | None = None
         self._cleanup_failure: BaseException | None = None
+        self._cleanup_truncated = False
         self._root: asyncio.Task[Any] | None = None
 
     def _now(self) -> datetime:
@@ -343,7 +343,7 @@ def resume(self) -> None:
     def _has_durable_work(self) -> bool:
         return any(not durable.is_complete for durable, _ in self._waits.values()) or bool(self._timers)
 
-    def _drain(self) -> int:
+    def _drain(self, cleanup_limit: int | None = None) -> int:
         previous_loop = asyncio._get_running_loop()
         previous_task = asyncio.current_task(previous_loop) if previous_loop is not None else None
         previous_hooks = sys.get_asyncgen_hooks()
@@ -353,7 +353,8 @@ def _drain(self) -> int:
         sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter, finalizer=self._asyncgen_finalizer)
         try:
             steps = 0
-            cleanup_limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
+            if cleanup_limit is None:
+                cleanup_limit = self._cleanup_budget()
             while self._ready or self._cancel_ready:
                 self._ready.extend(handle for _, _, handle in sorted(self._cancel_ready))
                 self._cancel_ready.clear()
@@ -379,10 +380,7 @@ def _drain(self) -> int:
                     self._deliver_completions()
                 steps += 1
                 if self._closing and steps >= cleanup_limit:
-                    self._record_cleanup_error(RuntimeError(
-                        "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
-                    self._ready.clear()
-                    self._cancel_ready.clear()
+                    self._report_unquiesced_cleanup()
                     break
             return steps
         finally:
@@ -411,12 +409,25 @@ def call_exception_handler(self, context: dict[str, Any]) -> None:
             error = failure
         self._fatal = error if isinstance(error, BaseException) else RuntimeError(context["message"])
 
+    def _cleanup_budget(self) -> int:
+        # Cooperative cleanup is capped in scheduling steps, not wall-clock time.
+        return max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
+
+    def _report_unquiesced_cleanup(self) -> None:
+        self._ready.clear()
+        self._cancel_ready.clear()
+        if not self._cleanup_truncated:
+            self._cleanup_truncated = True
+            self._record_cleanup_error(RuntimeError(
+                "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
+
     def _record_cleanup_error(self, error: BaseException, source: str = "cleanup callback") -> None:
         self._logger.warning("%s: Async %s failed.", self.ctx.instance_id, source, exc_info=error)
         if not isinstance(error, (Exception, asyncio.CancelledError)):
             if _is_process_control(self._cleanup_failure):
                 return
-            if self._cleanup_failure is None or is_control_flow_exception(error):
+            # Keep the first failure unless a process-control exit supersedes it.
+            if self._cleanup_failure is None or _is_process_control(error):
                 self._cleanup_failure = error
 
     def _close_asyncgens(self, closed: dict[int, Any]) -> None:
@@ -424,22 +435,24 @@ def _close_asyncgens(self, closed: dict[int, Any]) -> None:
         # Closed generators stay referenced so their ids cannot be reused during cleanup.
         steps = 0
         # Cleanup-created tasks and generators must not replenish the sweep's budget.
-        limit = max(10000, 4 * (len(self._tasks) + len(self._asyncgens)))
+        limit = self._cleanup_budget()
         while True:
             pending = [
                 generator for key, generator in self._asyncgens.items() if closed.get(key) is not generator]
             if not pending:
                 return
             if steps >= limit:
-                self._record_cleanup_error(RuntimeError(
-                    "Async cleanup did not quiesce; cleanup must cooperate with cancellation."))
+                self._report_unquiesced_cleanup()
                 return
             for generator in pending:
                 closed[id(generator)] = generator
                 self.create_task(generator.aclose())
-            steps += max(1, self._drain())
+            # Each drain spends from this sweep's remaining budget.
+            steps += max(1, self._drain(limit - steps))
 
     def close(self) -> None:
+        if self._closed:
+            return
         self._closing = True
         try:
             completed_at_close = {future for future in self._tasks if future.done()}
diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py
index bc821c8d..9bba5987 100644
--- a/tests/durabletask/test_asyncio_orchestrators.py
+++ b/tests/durabletask/test_asyncio_orchestrators.py
@@ -1126,6 +1126,60 @@ async def branch(index: int) -> None:
     assert cleaned == [0, 1]
 
 
+def test_cleanup_keeps_first_custom_exit_over_later_generator_exit() -> None:
+    class CleanupAbort(BaseException):
+        pass
+
+    errors: list[BaseException] = [CleanupAbort("first"), GeneratorExit("later")]
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def branch(index: int) -> None:
+            try:
+                await ctx.wait_for_external_event(f"never-{index}")
+            finally:
+                raise errors[index]
+
+        asyncio.create_task(branch(0))
+        asyncio.create_task(branch(1))
+        await ctx.call_activity("wait")
+
+    complete = _completion(_execute(orchestrator, [], _start()))
+    assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_FAILED
+    assert complete.failureDetails.errorType.endswith("CleanupAbort")
+
+
+def test_dispatch_restores_host_async_generator_hooks() -> None:
+    def firstiter(generator: Any) -> None:
+        pass
+
+    def finalizer(generator: Any) -> None:
+        pass
+
+    async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None:
+        async def resource() -> AsyncIterator[int]:
+            yield 1
+
+        await anext(resource())
+        await ctx.call_activity("wait")
+
+    previous = sys.get_asyncgen_hooks()
+    sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer)
+    try:
+        _execute(orchestrator, [], _start())
+        assert sys.get_asyncgen_hooks() == (firstiter, finalizer)
+    finally:
+        sys.set_asyncgen_hooks(*previous)
+
+
+def test_loop_close_is_idempotent(monkeypatch: pytest.MonkeyPatch) -> None:
+    ctx = cast(task.OrchestrationContext, SimpleNamespace(current_utc_datetime=NOW, instance_id="id"))
+    loop = DurableEventLoop(ctx, is_complete=lambda: False)
+    loop.close()
+    monkeypatch.setattr(loop, "_drain", lambda *args: pytest.fail("closed loop drained again"))
+    loop.close()
+    assert loop.is_closed()
+
+
 @pytest.mark.parametrize("count", [2, 32, 6000])
 def test_cleanup_closes_async_generators_started_during_generator_cleanup(
         count: int, caplog: pytest.LogCaptureFixture) -> None:
@@ -1146,10 +1200,13 @@ async def resource(index: int) -> AsyncIterator[int]:
 
     actions = _execute(orchestrator, [], _start()).actions
     assert [action.scheduleTask.name for action in actions] == ["wait"]
-    assert closed == list(range(count))
     if count == 6000:
-        assert "Async cleanup did not quiesce" in caplog.text
+        # Each sweep's step budget is fixed at entry, so a long chain is truncated.
+        assert caplog.text.count("Async cleanup did not quiesce") == 1
+        assert closed[:3000] == list(range(3000))
+        assert len(closed) < count
     else:
+        assert closed == list(range(count))
         assert "Async" not in caplog.text