From a2139cd5da752b6d69ddc3b295c8687251fbcbfb Mon Sep 17 00:00:00 2001 From: Andy Staples Date: Wed, 11 Nov 2026 13:12:16 -0700 Subject: [PATCH 1/5] Add cold async orchestration context Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .github/workflows/durabletask.yml | 1 + durabletask/async_context.py | 291 +++++++++++ durabletask/task.py | 3 +- durabletask/worker.py | 8 +- tests/durabletask/test_async_context.py | 491 ++++++++++++++++++ .../durabletask/test_asyncio_orchestrators.py | 66 +-- .../test_coroutine_orchestrators.py | 100 ++-- tests/durabletask/typing/async_context.py | 58 +++ 8 files changed, 943 insertions(+), 75 deletions(-) create mode 100644 durabletask/async_context.py create mode 100644 tests/durabletask/test_async_context.py create mode 100644 tests/durabletask/typing/async_context.py diff --git a/.github/workflows/durabletask.yml b/.github/workflows/durabletask.yml index 273f9e0a..b7272077 100644 --- a/.github/workflows/durabletask.yml +++ b/.github/workflows/durabletask.yml @@ -32,6 +32,7 @@ jobs: - name: Check asyncio, replay, and dispatch cleanup run: >- python -m pytest + tests\durabletask\test_async_context.py tests\durabletask\test_asyncio_orchestrators.py tests\durabletask\test_coroutine_orchestrators.py tests\durabletask\test_orchestration_generator_cleanup.py diff --git a/durabletask/async_context.py b/durabletask/async_context.py new file mode 100644 index 00000000..d8d293e7 --- /dev/null +++ b/durabletask/async_context.py @@ -0,0 +1,291 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Coroutine-native orchestration context backed by the durable replay engine.""" + +from __future__ import annotations + +import math +import logging +from datetime import datetime, timedelta +from typing import Any, TypeVar, overload + +from durabletask.entities import EntityInstanceId, EntityLock + +T = TypeVar("T") +TInput = TypeVar("TInput") +TOutput = TypeVar("TOutput") + + +class AsyncOrchestrationContext: + """Context supplied automatically to ``async def`` orchestrators. + + Durable operations return single-use native coroutines. Constructing one + schedules nothing: validation, input serialization, event consumption, and + relative deadline evaluation occur when it starts. Inputs are not copied + at construction; mutations before execution are observed. + + Cancelling activity, sub-orchestration, entity, or lock waiters does not + cancel remote work. Timers and event receives own their private local + registrations and withdraw them when their coroutine exits. This cannot + undo actions already committed to history or external side effects. + """ + + def __init__(self, context: task.OrchestrationContext) -> None: + self._context = context + + @property + def instance_id(self) -> str: + """Get the current orchestration instance ID.""" + return self._context.instance_id + + @property + def parent_instance_id(self) -> str | None: + """Get the parent orchestration instance ID, if any.""" + return self._context.parent_instance_id + + @property + def version(self) -> str | None: + """Get the orchestration instance version, if any.""" + return self._context.version + + @property + def current_utc_datetime(self) -> datetime: + """Get the replay-safe current UTC time.""" + return self._context.current_utc_datetime + + @property + def is_replaying(self) -> bool: + """Get whether the orchestration is replaying its history.""" + return self._context.is_replaying + + def set_custom_status(self, custom_status: Any) -> None: + """Set the orchestration's custom status.""" + self._context.set_custom_status(custom_status) + + def signal_entity( + self, entity_id: EntityInstanceId, operation_name: str, + input: Any = None, signal_time: datetime | None = None, + ) -> None: + """Send a one-way entity operation.""" + self._context.signal_entity(entity_id, operation_name, input, signal_time) + + def send_event(self, instance_id: str, event_name: str, *, data: Any | None = None) -> None: + """Send a one-way event to another orchestration.""" + self._context.send_event(instance_id, event_name, data=data) + + def continue_as_new( + self, new_input: Any, *, save_events: bool = False, new_version: str | None = None, + ) -> None: + """Continue the orchestration with new input.""" + self._context.continue_as_new(new_input, save_events=save_events, new_version=new_version) + + def new_uuid(self) -> str: + """Create a replay-safe UUID.""" + return self._context.new_uuid() + + def create_replay_safe_logger(self, logger: logging.Logger) -> task.ReplaySafeLogger: + """Create a logger that suppresses messages during replay.""" + return self._context.create_replay_safe_logger(logger) + + def as_task_context(self) -> task.OrchestrationContext: + """Return the eager task-handle context for advanced interoperability. + + Its operations schedule immediately and return reusable durable handles. + Await a handle directly or use ``handle.wait()`` for an independent + coroutine waiter whose cancellation does not cancel the underlying work. + """ + return self._context + + @overload + async def call_activity( + self, activity: task.Activity[TInput, TOutput] | str, *, + input: TInput | None = ..., + retry_policy: task.RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: type[T], + ) -> T: + ... + + @overload + async def call_activity( + self, activity: task.Activity[TInput, TOutput], *, + input: TInput | None = ..., + retry_policy: task.RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: None = ..., + ) -> TOutput: + ... + + @overload + async def call_activity( + self, activity: str, *, input: Any = ..., + retry_policy: task.RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: None = ..., + ) -> Any: + ... + + async def call_activity( + self, activity: task.Activity[TInput, TOutput] | str, *, + input: TInput | None = None, + retry_policy: task.RetryPolicy | None = None, + tags: dict[str, str] | None = None, + return_type: type | None = None, + ) -> Any: + """Start and await an activity when this coroutine executes.""" + return await self._context.call_activity( + activity, input=input, retry_policy=retry_policy, tags=tags, + return_type=return_type) + + @overload + async def call_sub_orchestrator( + self, orchestrator: task.Orchestrator[TInput, TOutput] | str, *, + input: TInput | None = ..., + instance_id: str | None = ..., + retry_policy: task.RetryPolicy | None = ..., + version: str | None = ..., + return_type: type[T], + ) -> T: + ... + + @overload + async def call_sub_orchestrator( + self, orchestrator: task.Orchestrator[TInput, TOutput], *, + input: TInput | None = ..., + instance_id: str | None = ..., + retry_policy: task.RetryPolicy | None = ..., + version: str | None = ..., + return_type: None = ..., + ) -> TOutput: + ... + + @overload + async def call_sub_orchestrator( + self, orchestrator: str, *, input: Any = ..., + instance_id: str | None = ..., + retry_policy: task.RetryPolicy | None = ..., + version: str | None = ..., + return_type: None = ..., + ) -> Any: + ... + + async def call_sub_orchestrator( + self, orchestrator: task.Orchestrator[TInput, TOutput] | str, *, + input: TInput | None = None, + instance_id: str | None = None, + retry_policy: task.RetryPolicy | None = None, + version: str | None = None, + return_type: type | None = None, + ) -> Any: + """Start and await a sub-orchestration when this coroutine executes.""" + return await self._context.call_sub_orchestrator( + orchestrator, input=input, instance_id=instance_id, + retry_policy=retry_policy, version=version, return_type=return_type) + + @overload + async def call_entity( + self, entity: EntityInstanceId, operation: str, input: Any = ..., *, + return_type: type[T], + ) -> T: + ... + + @overload + async def call_entity( + self, entity: EntityInstanceId, operation: str, input: Any = ..., *, + return_type: None = ..., + ) -> Any: + ... + + async def call_entity( + self, entity: EntityInstanceId, operation: str, input: Any = None, *, + return_type: type | None = None, + ) -> Any: + """Start and await an entity operation when this coroutine executes.""" + return await self._context.call_entity( + entity, operation, input, return_type=return_type) + + async def lock_entities(self, entities: list[EntityInstanceId]) -> EntityLock: + """Acquire locks, for use as ``with await ctx.lock_entities(ids):``. + + The returned synchronous context manager releases the locks. This + preserves the original context's entity locking and replay rules; + cancellation does not withdraw an already-issued lock request. An + abandoned request can still acquire locks, which remain held until + released or the orchestration ends. Retain and shield the acquisition + task if a timeout or another local cancellation might interrupt it, + then await it and exit the returned lock when it is safe to do so. + """ + return await self._context.lock_entities(entities) + + async def create_timer(self, fire_at: datetime | timedelta) -> None: + """Await a private durable timer, cancelling it when this wait exits. + + Relative deadlines are evaluated when execution starts, not when this + coroutine is constructed. Cancellation removes uncommitted timer work + where possible; it cannot retract a timer already present in history. + """ + timer = self._context.create_timer(fire_at) + try: + await timer + finally: + timer.cancel() + + @overload + async def wait_for_external_event( + self, name: str, *, timeout: timedelta | None = ..., + data_type: type[T], + ) -> T: + ... + + @overload + async def wait_for_external_event( + self, name: str, *, timeout: timedelta | None = ..., + data_type: None = ..., + ) -> Any: + ... + + async def wait_for_external_event( + self, name: str, *, timeout: timedelta | None = None, + data_type: type | None = None, + ) -> Any: + """Receive one event through a private, cancellation-owned subscription. + + A finite relative timeout starts when the coroutine executes and raises + ``TimeoutError`` if its durable timer wins. Buffered events win + immediately, even for zero/negative timeouts, without creating a timer. + Otherwise nonpositive timeouts fail immediately without a timer. + Pending races follow history order. On every exit, the subscription + and any outstanding timeout timer are withdrawn locally, so a later + receive cannot lose an event to this abandoned subscription. + + For a shared subscription use ``as_task_context().wait_for_external_event()`` + and borrowed ``Task.wait()`` instead. + """ + from durabletask.task import when_any + + if timeout is not None and not math.isfinite(timeout.total_seconds()): + raise ValueError("Event receive timeouts must be finite.") + event: task.CancellableTask[Any] = self._context.wait_for_external_event(name, data_type=data_type) + timer: task.TimerTask | None = None + try: + if event.is_complete: + return event.get_result() + if timeout is None: + return await event + if timeout <= timedelta(0): + raise TimeoutError(f"Timed out waiting for event {name!r}.") + timer = self._context.create_timer(timeout) + winner = await when_any([event, timer]) + if winner is timer: + raise TimeoutError(f"Timed out waiting for event {name!r}.") + return event.get_result() + finally: + event.cancel() + if timer is not None: + timer.cancel() + + +# Resolve public annotations at runtime without importing task before this +# module has defined the context that task re-exports. +from durabletask import task # noqa: E402 diff --git a/durabletask/task.py b/durabletask/task.py index 483b4767..73578010 100644 --- a/durabletask/task.py +++ b/durabletask/task.py @@ -13,6 +13,7 @@ from datetime import UTC, datetime, timedelta from typing import Any, Generic, TypeAlias, TypeVar, cast, overload +from durabletask.async_context import AsyncOrchestrationContext as AsyncOrchestrationContext from durabletask.entities import DurableEntity, EntityInstanceId, EntityLock, EntityContext import durabletask.internal.helpers as pbh import durabletask.internal.orchestrator_service_pb2 as pb @@ -986,7 +987,7 @@ def task_id(self) -> int: ] AsyncOrchestrator: TypeAlias = Callable[ - [OrchestrationContext, TInput], Coroutine[Any, Any, TOutput], + [AsyncOrchestrationContext, TInput], Coroutine[Any, Any, TOutput], ] Orchestrator: TypeAlias = SyncOrchestrator[TInput, TOutput] | AsyncOrchestrator[TInput, TOutput] diff --git a/durabletask/worker.py b/durabletask/worker.py index df31c42c..1b6ccaec 100644 --- a/durabletask/worker.py +++ b/durabletask/worker.py @@ -2744,9 +2744,13 @@ def process_event( input = self._data_converter.deserialize( event.executionStarted.input.value, input_type) - # Layer 4 passes async orchestrators a cold context instead. is_async = shared.is_async_callable(fn) - result = fn(ctx, input) + if is_async: + async_fn = cast(task.AsyncOrchestrator[Any, Any], fn) + result = async_fn(task.AsyncOrchestrationContext(ctx), input) + else: + sync_fn = cast(task.SyncOrchestrator[Any, Any], fn) + result = sync_fn(ctx, input) if isinstance(result, CoroutineType) and not is_async: result.close() diff --git a/tests/durabletask/test_async_context.py b/tests/durabletask/test_async_context.py new file mode 100644 index 00000000..d822f2c2 --- /dev/null +++ b/tests/durabletask/test_async_context.py @@ -0,0 +1,491 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Cold coroutine and ownership semantics through the shared replay executor.""" + +import asyncio +import inspect +import json +import logging +from collections.abc import Coroutine, Generator +from datetime import timedelta +from typing import Any, get_type_hints + +import pytest + +from durabletask import entities, task +from durabletask.async_context import AsyncOrchestrationContext +from durabletask.internal import helpers +from durabletask.internal import orchestrator_service_pb2 as pb +from tests.durabletask.test_coroutine_orchestrators import ( + NOW, + _completion, # pyright: ignore[reportPrivateUsage] + _execute, # pyright: ignore[reportPrivateUsage] + _start, # pyright: ignore[reportPrivateUsage] +) + + +def test_additive_public_api_has_no_new_abstract_requirements() -> None: + assert "as_task_context" not in task.OrchestrationContext.__abstractmethods__ + assert "wait" not in task.Task.__abstractmethods__ + assert task.AsyncOrchestrationContext is AsyncOrchestrationContext + + +def test_public_async_context_annotations_resolve_at_runtime() -> None: + assert get_type_hints(AsyncOrchestrationContext.call_activity)["return"] is Any + assert get_type_hints(AsyncOrchestrationContext.as_task_context)["return"] is task.OrchestrationContext + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_all_operations_are_cold_and_prestart_cancellation_schedules_nothing(cancel: bool) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: + ops = ctx + entity = entities.EntityInstanceId("counter", "one") + invalid_input = object() + operations: list[Coroutine[Any, Any, Any]] = [ + ops.call_activity("unused", input=invalid_input), + ops.call_sub_orchestrator("unused", input=invalid_input), + ops.call_entity(entity, "unused", invalid_input), + ops.lock_entities([entity]), + ops.create_timer(timedelta(days=1)), + ops.wait_for_external_event("unused", timeout=timedelta(days=1)), + ] + assert all(inspect.iscoroutine(operation) for operation in operations) + if cancel: + pending = [asyncio.create_task(operation) for operation in operations] + for waiter in pending: + waiter.cancel() + await asyncio.gather(*pending, return_exceptions=True) + else: + for operation in operations: + operation.close() + return await ctx.call_activity("actual", input=3) + + actions = _execute(orchestrator, [], _start()).actions + assert len(actions) == 1 + assert actions[0].id == 1 + assert actions[0].scheduleTask.name == "actual" + + +@pytest.mark.parametrize("cancel", [False, True]) +def test_cold_receive_does_not_consume_buffered_event(cancel: bool) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: + await ctx.wait_for_external_event("go") + unused = ctx.wait_for_external_event("value") + if cancel: + pending = asyncio.create_task(unused) + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + else: + unused.close() + return await ctx.wait_for_external_event("value", data_type=str) + + result = _execute(orchestrator, _start(), [ + helpers.new_event_raised_event("value", '"kept"'), + helpers.new_event_raised_event("go", None), + ]) + assert json.loads(_completion(result).result.value) == "kept" + + +def test_create_task_and_task_group_accept_cold_operations() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[int]: + ops = ctx + first = asyncio.create_task(ops.call_activity("first", return_type=int)) + async with asyncio.TaskGroup() as group: + second = group.create_task(ops.call_activity("second", return_type=int)) + return [await first, second.result()] + + start = _start() + initial = _execute(orchestrator, [], start) + assert [(a.id, a.scheduleTask.name) for a in initial.actions] == [(1, "first"), (2, "second")] + old = start + [ + helpers.new_task_scheduled_event(1, "first"), + helpers.new_task_scheduled_event(2, "second"), + ] + result = _execute(orchestrator, old, [ + helpers.new_task_completed_event(2, "2"), + helpers.new_task_completed_event(1, "1"), + ]) + assert json.loads(_completion(result).result.value) == [1, 2] + + +def test_serialization_and_scheduling_follow_coroutine_execution_not_construction() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: + ops = ctx + payload = {"value": 1} + first = ops.call_activity("first", input=payload) + second = ops.call_activity("second", input=payload) + payload["value"] = 2 + try: + await second + return await first + finally: + first.close() + + actions = _execute(orchestrator, [], _start()).actions + assert len(actions) == 1 + assert actions[0].id == 1 + assert actions[0].scheduleTask.name == "second" + assert json.loads(actions[0].scheduleTask.input.value) == {"value": 2} + + +def test_cold_relative_timer_deadline_is_evaluated_when_started() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: + timer = ctx.create_timer(timedelta(hours=1)) + try: + await ctx.wait_for_external_event("go") + await timer + finally: + timer.close() + + later = NOW + timedelta(hours=2) + result = _execute(orchestrator, _start(), [ + helpers.new_orchestrator_started_event(later), + helpers.new_event_raised_event("go", None), + ]) + assert len(result.actions) == 1 + assert result.actions[0].createTimer.fireAt.ToDatetime() == later + timedelta(hours=1) + + +def test_borrowed_waiters_are_independent_reusable_and_do_not_cancel_operations() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[str]: + operation = ctx.as_task_context().call_activity("shared", return_type=str) + cancelled = asyncio.create_task(operation.wait()) + remaining = asyncio.create_task(operation.wait()) + await asyncio.sleep(0) + cancelled.cancel() + await asyncio.gather(cancelled, return_exceptions=True) + return [await remaining, await operation.wait(), await operation] + + old = _start() + [helpers.new_task_scheduled_event(1, "shared")] + result = _execute(orchestrator, old, [helpers.new_task_completed_event(1, '"ok"')]) + assert json.loads(_completion(result).result.value) == ["ok", "ok", "ok"] + + +def test_borrowed_event_cancellation_preserves_shared_subscription() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + event = ctx.as_task_context().wait_for_external_event("value", data_type=int) + first = asyncio.create_task(event.wait()) + second = asyncio.create_task(event.wait()) + await asyncio.sleep(0) + first.cancel() + await asyncio.gather(first, return_exceptions=True) + return await second + + result = _execute(orchestrator, _start(), [helpers.new_event_raised_event("value", "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +def test_borrowed_cancelled_subscription_still_receives_before_owned_replacement() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[int]: + shared = ctx.as_task_context().wait_for_external_event("value", data_type=int) + pending = asyncio.create_task(shared.wait()) + await asyncio.sleep(0) + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + next_value = await ctx.wait_for_external_event("value", data_type=int) + return [await shared.wait(), next_value] + + result = _execute(orchestrator, _start(), [ + helpers.new_event_raised_event("value", "1"), + helpers.new_event_raised_event("value", "2"), + ]) + assert json.loads(_completion(result).result.value) == [1, 2] + + +def test_cold_coroutines_are_single_use_not_shared_operation_handles() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + operation = ctx.call_activity("work", return_type=int) + value = await operation + with pytest.raises(RuntimeError, match="cannot reuse"): + await operation + return value + + old = _start() + [helpers.new_task_scheduled_event(1, "work")] + result = _execute(orchestrator, old, [helpers.new_task_completed_event(1, "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +def test_owned_timeout_withdraws_subscription_before_replacement() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + ops = ctx + try: + await ops.wait_for_external_event("value", timeout=timedelta(seconds=5)) + except TimeoutError: + return await ops.wait_for_external_event("value", data_type=int) + raise AssertionError("Expected timeout") + + deadline = NOW + timedelta(seconds=5) + initial = _execute(orchestrator, [], _start()) + assert initial.actions[0].id == 1 + old = _start() + [helpers.new_timer_created_event(1, deadline)] + expired = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(1, deadline)] + assert _execute(orchestrator, old, expired).actions == [] + result = _execute(orchestrator, old + expired, [helpers.new_event_raised_event("value", "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +@pytest.mark.parametrize("with_timeout", [False, True]) +def test_owned_cancellation_withdraws_subscription_and_uncommitted_timer(with_timeout: bool) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + ops = ctx + pending = asyncio.create_task(ops.wait_for_external_event( + "value", timeout=timedelta(seconds=5) if with_timeout else None)) + await asyncio.sleep(0) + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + return await ops.wait_for_external_event("value", data_type=int) + + assert _execute(orchestrator, [], _start()).actions == [] + result = _execute(orchestrator, _start(), [helpers.new_event_raised_event("value", "42")]) + assert json.loads(_completion(result).result.value) == 42 + + +@pytest.mark.parametrize("buffered", [False, True]) +@pytest.mark.parametrize("seconds", [-1, 0, 5]) +def test_buffered_events_and_immediate_timeouts(buffered: bool, seconds: int) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: + await ctx.wait_for_external_event("go") + try: + return await ctx.wait_for_external_event( + "value", timeout=timedelta(seconds=seconds), data_type=int) + except TimeoutError: + return "timeout" + + new = ([helpers.new_event_raised_event("value", "42")] if buffered else []) + [ + helpers.new_event_raised_event("go", None)] + result = _execute(orchestrator, _start(), new) + if not buffered and seconds > 0: + assert len(result.actions) == 1 + assert result.actions[0].HasField("createTimer") + else: + assert len(result.actions) == 1 + assert json.loads(_completion(result).result.value) == (42 if buffered else "timeout") + + +@pytest.mark.parametrize("event_first", [False, True]) +def test_owned_receive_race_follows_history_order(event_first: bool) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: + ops = ctx + try: + value = await ops.wait_for_external_event("value", timeout=timedelta(seconds=5), data_type=int) + return f"event:{value}" + except TimeoutError: + return f"timeout:{await ops.wait_for_external_event('value', data_type=int)}" + + deadline = NOW + timedelta(seconds=5) + old = _start() + [helpers.new_timer_created_event(1, deadline)] + arrival = helpers.new_event_raised_event("value", "42") + fired = helpers.new_timer_fired_event(1, deadline) + result = _execute(orchestrator, old, [arrival, fired] if event_first else [fired, arrival]) + assert json.loads(_completion(result).result.value) == ("event:42" if event_first else "timeout:42") + + +@pytest.mark.parametrize("seconds", [float("inf"), float("-inf"), float("nan")]) +def test_nonfinite_timeout_rejected_before_consuming_event(seconds: float) -> None: + class InfiniteTimeout(timedelta): + def total_seconds(self) -> float: + return seconds + + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + await ctx.wait_for_external_event("go") + ops = ctx + with pytest.raises(ValueError, match="finite"): + await ops.wait_for_external_event("value", timeout=InfiniteTimeout()) + return await ops.wait_for_external_event("value", data_type=int) + + result = _execute(orchestrator, _start(), [ + helpers.new_event_raised_event("value", "42"), + helpers.new_event_raised_event("go", None), + ]) + assert json.loads(_completion(result).result.value) == 42 + + +def test_owned_timer_cancellation_does_not_cancel_borrowed_timer() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: + borrowed = ctx.as_task_context().create_timer(timedelta(seconds=10)) + owned = asyncio.create_task(ctx.create_timer(timedelta(seconds=5))) + waiter = asyncio.create_task(borrowed.wait()) + await asyncio.sleep(0) + owned.cancel() + waiter.cancel() + await asyncio.gather(owned, waiter, return_exceptions=True) + await borrowed + + result = _execute(orchestrator, [], _start()) + assert len(result.actions) == 1 + assert result.actions[0].id == 1 + assert result.actions[0].createTimer.fireAt.ToDatetime() == NOW + timedelta(seconds=10) + + +def test_owned_timer_completes_through_history() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: + await asyncio.create_task(ctx.create_timer(timedelta(seconds=5))) + return "done" + + deadline = NOW + timedelta(seconds=5) + old = _start() + [helpers.new_timer_created_event(1, deadline)] + result = _execute(orchestrator, old, [helpers.new_timer_fired_event(1, deadline)]) + assert json.loads(_completion(result).result.value) == "done" + + +def test_owned_event_completion_cancels_only_its_uncommitted_timeout_timer() -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: + await ctx.wait_for_external_event("value", timeout=timedelta(seconds=5)) + return await ctx.call_activity("next") + + result = _execute(orchestrator, _start(), [helpers.new_event_raised_event("value", "42")]) + assert len(result.actions) == 1 + assert result.actions[0].id == 2 + assert result.actions[0].scheduleTask.name == "next" + + +def test_dispatch_disposal_of_owned_resources_preserves_response_and_has_no_cleanup_errors( + caplog: pytest.LogCaptureFixture, +) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: + ops = ctx + await asyncio.gather( + ops.create_timer(timedelta(seconds=10)), + ops.wait_for_external_event("value", timeout=timedelta(seconds=5)), + ops.wait_for_external_event("other"), + ) + + with caplog.at_level(logging.WARNING): + result = _execute(orchestrator, [], _start()) + assert [(a.id, a.createTimer.fireAt.ToDatetime()) for a in result.actions] == [ + (1, NOW + timedelta(seconds=10)), (2, NOW + timedelta(seconds=5))] + assert not caplog.records + + +def test_sequential_cold_operations_preserve_generator_actions() -> None: + async def coroutine(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + ops = ctx + first = await ops.call_activity("increment", input=3, return_type=int) + return await ops.call_activity("increment", input=first, return_type=int) + + def generator(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, int]: + first = yield ctx.call_activity("increment", input=3, 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, "4")], + [helpers.new_task_scheduled_event(2, "increment"), helpers.new_task_completed_event(2, "5")], + ] + old: list[pb.HistoryEvent] = [] + result = 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_entity_lock_and_typed_call_use_existing_replay_protocol() -> None: + entity = entities.EntityInstanceId("counter", "one") + + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + ops = ctx + with await ops.lock_entities([entity]): + return await ops.call_entity(entity, "get", return_type=int) + + old = _start() + lock = _execute(orchestrator, [], old).actions[0].sendEntityMessage.entityLockRequested + new = [ + helpers.new_event_sent_event(1, str(entity), json.dumps({"id": lock.criticalSectionId})), + helpers.new_event_raised_event(lock.criticalSectionId, None), + ] + actions = _execute(orchestrator, old, new).actions + call = actions[0].sendEntityMessage.entityOperationCalled + 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_timed_out_shielded_lock_is_released_after_acquisition() -> None: + entity = entities.EntityInstanceId("counter", "one") + + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: + acquire = asyncio.create_task(ctx.lock_entities([entity])) + try: + lock = await asyncio.wait_for(asyncio.shield(acquire), timeout=2) + except TimeoutError: + lock = await acquire + with lock: + pass + return "timed out" + with lock: + return "acquired" + + start = _start() + initial = _execute(orchestrator, [], start) + assert len(initial.actions) == 2 + lock_action = next(action for action in initial.actions if action.HasField("sendEntityMessage")) + timer_action = next(action for action in initial.actions if action.HasField("createTimer")) + section = lock_action.sendEntityMessage.entityLockRequested.criticalSectionId + deadline = timer_action.createTimer.fireAt.ToDatetime() + assert deadline == NOW + timedelta(seconds=2) + history = start + [ + helpers.new_event_sent_event(lock_action.id, str(entity), json.dumps({"id": section})), + helpers.new_timer_created_event(timer_action.id, deadline), + ] + timed_out = [ + helpers.new_orchestrator_started_event(deadline), + helpers.new_timer_fired_event(timer_action.id, deadline), + ] + assert not _execute(orchestrator, history, timed_out).actions + history += timed_out + granted = [helpers.new_event_raised_event(section, None)] + result = _execute(orchestrator, history, granted) + assert result.actions == _execute(orchestrator, history, granted).actions + assert result.actions[0].sendEntityMessage.entityUnlockSent.targetInstanceId.value == str(entity) + assert json.loads(_completion(result).result.value) == "timed out" + + +@pytest.mark.parametrize("sub_orchestration", [False, True]) +def test_cold_retry_and_call_options_preserve_eager_actions(sub_orchestration: bool) -> None: + policy = task.RetryPolicy(first_retry_interval=timedelta(seconds=2), max_number_of_attempts=3) + + async def cold(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + if sub_orchestration: + return await ctx.call_sub_orchestrator( + "child", input=3, instance_id="child-id", retry_policy=policy, + version="v1", return_type=int) + return await ctx.call_activity( + "work", input=3, retry_policy=policy, tags={"tag": "value"}, return_type=int) + + def eager(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, int]: + if sub_orchestration: + return (yield ctx.call_sub_orchestrator( + "child", input=3, instance_id="child-id", retry_policy=policy, + version="v1", return_type=int)) + return (yield ctx.call_activity( + "work", input=3, retry_policy=policy, tags={"tag": "value"}, return_type=int)) + + start = _start() + assert _execute(cold, [], start).actions == _execute(eager, [], start).actions + if sub_orchestration: + scheduled = helpers.new_sub_orchestration_created_event(1, "child", "child-id") + failed = helpers.new_sub_orchestration_failed_event(1, ValueError("retry me")) + completed = helpers.new_sub_orchestration_completed_event(1, "42") + else: + scheduled = helpers.new_task_scheduled_event(1, "work") + failed = helpers.new_task_failed_event(1, ValueError("retry me")) + completed = helpers.new_task_completed_event(1, "42") + old = start + [scheduled] + retry = _execute(cold, old, [failed]) + assert retry.actions == _execute(eager, old, [failed]).actions + assert len(retry.actions) == 1 + assert retry.actions[0].id == 2 + deadline = retry.actions[0].createTimer.fireAt.ToDatetime() + old += [failed, helpers.new_timer_created_event(2, deadline)] + fired = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(2, deadline)] + restarted = _execute(cold, old, fired) + assert restarted.actions == _execute(eager, old, fired).actions + assert restarted.actions[0].id == 1 + old += fired + [scheduled] + complete = _completion(_execute(cold, old, [completed])) + assert json.loads(complete.result.value) == 42 diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py index 3fcf1b9b..a55ec596 100644 --- a/tests/durabletask/test_asyncio_orchestrators.py +++ b/tests/durabletask/test_asyncio_orchestrators.py @@ -34,9 +34,9 @@ @pytest.mark.parametrize("fail_branch", [False, True]) def test_task_group_branches_inherit_context_and_replay_stable_actions(fail_branch: bool) -> None: - context: ContextVar[task.OrchestrationContext] = ContextVar("orchestration") + context: ContextVar[task.AsyncOrchestrationContext] = ContextVar("orchestration") - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: token = context.set(ctx) try: results: dict[str, Any] = {} @@ -86,7 +86,7 @@ def replay(old: list[pb.HistoryEvent], new: list[pb.HistoryEvent]): def test_gather_resumes_each_branch_without_waiting_for_batch() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[Any]: async def branch(name: str) -> Any: value: Any = await ctx.call_activity(name) await asyncio.sleep(0) @@ -112,12 +112,12 @@ async def branch(name: str) -> Any: @pytest.mark.parametrize("return_exceptions", [False, True]) def test_gather_failure_does_not_cancel_sibling(return_exceptions: bool) -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: async def sibling() -> str: await ctx.call_activity("slow") return await ctx.call_activity("continued", return_type=str) - failure = asyncio.create_task(ctx.call_activity("fail", return_type=str).wait()) + failure = asyncio.create_task(ctx.call_activity("fail", return_type=str)) pending = asyncio.create_task(sibling()) try: values = await asyncio.gather( @@ -141,7 +141,7 @@ async def sibling() -> str: def test_task_group_failure_cancels_waiters_not_durable_work() -> None: cleaned: list[str] = [] - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def branch(name: str) -> None: try: await ctx.call_activity(name) @@ -162,7 +162,7 @@ async def branch(name: str) -> None: @pytest.mark.parametrize("kind", ["sleep", "wait_for", "timeout"]) def test_durable_timer_and_timeout(kind: str) -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: try: if kind == "sleep": await asyncio.sleep(2) @@ -186,7 +186,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: def test_sleep_zero_fifo_without_timer() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[int]: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[int]: observed: list[int] = [] async def branch(index: int) -> None: @@ -212,7 +212,7 @@ class Payload: def __post_init__(self) -> None: references.append(ref(self)) - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: for batch in range(8): await asyncio.gather(*(ctx.call_activity("work", return_type=Payload) for _ in range(4))) await asyncio.sleep(0) @@ -248,7 +248,7 @@ class Payload: def __post_init__(self) -> None: references.append(ref(self)) - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: async def branch() -> None: payload = await ctx.call_activity("work", return_type=Payload) raise ValueError(f"Failed with payload {payload.value}") @@ -308,7 +308,7 @@ async def branch() -> None: def test_retired_tasks_do_not_reuse_cancellation_order() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: async def completed() -> None: pass @@ -342,7 +342,7 @@ def test_dispatch_cleanup_unwinds_scopes_contexts_asyncgens(under_loop: bool) -> cleaned: list[str] = [] references: list[Any] = [] - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: references.extend([ref(ctx), ref(asyncio.get_running_loop()), ref(asyncio.current_task())]) async def resource(): @@ -381,8 +381,8 @@ async def outer() -> None: 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) + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: + durable = ctx.as_task_context().call_activity("shared", return_type=str) async def wait() -> Any: return await durable @@ -400,8 +400,8 @@ async def wait() -> Any: def test_local_durable_cancellation_wakes_waiter_without_history() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: - timer = ctx.create_timer(timedelta(days=1)) + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: + timer = ctx.as_task_context().create_timer(timedelta(days=1)) async def waiter() -> None: await timer @@ -433,7 +433,7 @@ async def resource(): await asyncio.sleep(0) cleaned.append("resource") - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async with resource(): try: ctx.set_custom_status("waiting") @@ -472,7 +472,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: def test_unhandled_local_root_cancellation_fails_not_abandons() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: raise asyncio.CancelledError complete = _completion(_execute(orchestrator, [], _start())) @@ -484,7 +484,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: def test_timer_rounding_never_precedes_loop_deadline(delay: float) -> None: observed: list[float] = [] - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: loop = asyncio.get_running_loop() observed.append(loop.time() + delay) await asyncio.sleep(delay) @@ -556,7 +556,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: def test_async_sleep_uses_shared_long_timer_chunking() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: await asyncio.sleep(7 * 86400) return "done" @@ -576,7 +576,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: @pytest.mark.parametrize("api", ["thread", "io", "foreign-loop"]) def test_unsupported_work_fails_actionably(api: str) -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: if api == "thread": await asyncio.to_thread(lambda: None) elif api == "io": @@ -603,7 +603,7 @@ class CleanupAbort(BaseException): error = CleanupAbort("cleanup failed") if unexpected else ValueError("cleanup failed") - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def resource(): try: yield 1 @@ -631,7 +631,7 @@ async def resource(): def test_task_group_cancellation_durable_cleanup_has_stable_action_order() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def branch(index: int) -> None: try: await ctx.call_activity("work", input=index) @@ -768,7 +768,7 @@ def test_shared_future_cancellation_history_replays_in_other_worker_processes() def test_cancelling_last_loop_timer_cannot_hide_foreign_wait() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: handle = asyncio.get_running_loop().call_later(1, lambda: None) handle.cancel() await asyncio.Future() @@ -782,7 +782,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: @pytest.mark.parametrize("probe_pending", [False, True]) def test_background_failure_is_reported_unless_observed( observed: bool, probe_pending: bool, caplog: pytest.LogCaptureFixture) -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: async def fail() -> None: raise ValueError("background failure") @@ -809,7 +809,7 @@ def test_unexpected_cleanup_failure_does_not_hide_other_task_errors( class CleanupAbort(BaseException): pass - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def branch(error: BaseException) -> None: try: await ctx.wait_for_external_event("never") @@ -835,7 +835,7 @@ def test_cleanup_callback_failure_still_drains_tasks_and_asyncgens( def fail() -> None: raise SystemExit("cleanup exit") if control_flow else ValueError("cleanup callback") - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def resource(): try: yield 1 @@ -890,7 +890,7 @@ def is_complete(durable: task.Task[Any]) -> bool: monkeypatch.setattr(task.Task, "is_complete", property(is_complete)) - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[Any]: async def branch(index: int) -> Any: return await ctx.call_activity("work", input=index) @@ -906,7 +906,7 @@ 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 orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def resource(): try: yield 1 @@ -939,7 +939,7 @@ class CleanupAbort(BaseException): if grouped: exit_error = BaseExceptionGroup("forced cleanup group", [exit_error, ValueError("nested")]) - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def branch(name: str, error: BaseException | None) -> None: try: await ctx.wait_for_external_event(name) @@ -970,7 +970,7 @@ async def branch(name: str, error: BaseException | None) -> None: @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: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: def fail() -> None: raise StopIteration("callback stopped") @@ -997,7 +997,7 @@ def fail() -> None: @pytest.mark.parametrize("cancelled", [None, 0, 1]) def test_identical_timer_callbacks_complete_and_cancel_independently(cancelled: int | None) -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> int: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: loop = asyncio.get_running_loop() result = loop.create_future() completed: list[None] = [] @@ -1039,7 +1039,7 @@ class CleanupAbort(BaseException): exit_error = BaseExceptionGroup("cleanup exit", [SystemExit("stop"), ValueError("nested")]) cleaned: list[str] = [] - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def branch(name: str, error: BaseException | None) -> None: try: await ctx.wait_for_external_event(name) diff --git a/tests/durabletask/test_coroutine_orchestrators.py b/tests/durabletask/test_coroutine_orchestrators.py index 824b3afe..9c2d5dbb 100644 --- a/tests/durabletask/test_coroutine_orchestrators.py +++ b/tests/durabletask/test_coroutine_orchestrators.py @@ -52,7 +52,7 @@ def _completion(result: worker.ExecutionResults) -> pb.CompleteOrchestrationActi def test_coroutine_immediate_return() -> None: - async def orchestrator(ctx: task.OrchestrationContext, value: int) -> int: + async def orchestrator(ctx: task.AsyncOrchestrationContext, value: int) -> int: return value complete = _completion(_execute(orchestrator, [], _start())) @@ -61,25 +61,29 @@ async def orchestrator(ctx: task.OrchestrationContext, value: int) -> int: @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: +def test_async_context_is_selected_for_supported_callable_shapes(shape: str) -> None: + async def plain(ctx: task.AsyncOrchestrationContext, value: int) -> int: + assert isinstance(ctx, task.AsyncOrchestrationContext) 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: + async def with_prefix(prefix: int, ctx: task.AsyncOrchestrationContext, value: int) -> int: + assert isinstance(ctx, task.AsyncOrchestrationContext) return value + prefix class AsyncCallable: - async def __call__(self, ctx: task.OrchestrationContext, value: int) -> int: + async def __call__(self, ctx: task.AsyncOrchestrationContext, value: int) -> int: + assert isinstance(ctx, task.AsyncOrchestrationContext) return value + 1 class AsyncPartialCallable: async def __call__( - self, prefix: int, ctx: task.OrchestrationContext, value: int, + self, prefix: int, ctx: task.AsyncOrchestrationContext, value: int, ) -> int: + assert isinstance(ctx, task.AsyncOrchestrationContext) return value + prefix variants = { @@ -166,7 +170,7 @@ def test_synchronous_builtin_callables_are_not_classified_as_async(fn: Any) -> N 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: + async def async_orchestrator(ctx: task.AsyncOrchestrationContext, value: int) -> int: return value def wrapper(ctx: task.OrchestrationContext, value: int) -> Any: @@ -181,7 +185,7 @@ def wrapper(ctx: task.OrchestrationContext, value: int) -> Any: def test_async_classified_wrapper_returning_non_coroutine_is_rejected() -> None: - async def async_orchestrator(ctx: task.OrchestrationContext, value: int) -> int: + async def async_orchestrator(ctx: task.AsyncOrchestrationContext, value: int) -> int: return value @wraps(async_orchestrator) @@ -196,7 +200,7 @@ def wrapper(*args: Any, **kwargs: Any) -> int: 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: + async def async_orchestrator(ctx: task.AsyncOrchestrationContext, value: int) -> int: return value @wraps(async_orchestrator) @@ -262,10 +266,41 @@ def test_mixed_wrapper_cycle_fails_within_bounded_worker_dispatch(shape: str) -> assert result.returncode == 0, result.stdout + result.stderr +def test_async_context_metadata_and_commands_match_sync_context() -> None: + entity = entities.EntityInstanceId("counter", "one") + + def use_context(ctx: task.OrchestrationContext | task.AsyncOrchestrationContext) -> list[Any]: + ctx.set_custom_status({"state": "running"}) + ctx.send_event("other", "ping", data=3) + ctx.signal_entity(entity, "add", 1) + return [ + ctx.instance_id, ctx.parent_instance_id, ctx.version, + ctx.current_utc_datetime.isoformat(), ctx.is_replaying, ctx.new_uuid(), + ] + + async def async_orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[Any]: + return use_context(ctx) + + def sync_orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: + return use_context(ctx) + + started = helpers.new_execution_started_event("test", INSTANCE_ID, "null", version="v1") + started.executionStarted.parentInstance.CopyFrom(pb.ParentInstanceInfo( + taskScheduledId=1, + orchestrationInstance=pb.OrchestrationInstance(instanceId="parent"))) + new = [helpers.new_orchestrator_started_event(NOW), started] + async_result = _execute(async_orchestrator, [], new) + sync_result = _execute(sync_orchestrator, [], new) + assert async_result.actions == sync_result.actions + assert async_result.encoded_custom_status == sync_result.encoded_custom_status == '{"state": "running"}' + assert json.loads(_completion(async_result).result.value)[:5] == [ + INSTANCE_ID, "parent", "v1", NOW.isoformat(), False] + + 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]: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[int]: completed = task.CompletableTask[int]() completed.complete(7) return [await completed, await completed] @@ -273,19 +308,6 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[int]: 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]: @@ -303,7 +325,7 @@ def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Generator[Any, Any, def test_precompleted_failure_and_cancellation_are_catchable() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[str]: failed = task.CompletableTask[int]() failed.fail("failed", ValueError("boom")) cancelled = task.TimerTask() @@ -324,10 +346,10 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]: def test_nested_async_helpers_replay_with_generator_action_parity() -> None: - async def helper(ctx: task.OrchestrationContext, value: int) -> int: + async def helper(ctx: task.AsyncOrchestrationContext, value: int) -> int: return await ctx.call_activity("increment", input=value, return_type=int) - async def coroutine(ctx: task.OrchestrationContext, value: int) -> int: + async def coroutine(ctx: task.AsyncOrchestrationContext, value: int) -> int: first = await helper(ctx, value) return await helper(ctx, first) @@ -353,9 +375,9 @@ def generator(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task def test_fan_out_waits_for_all_and_preserves_input_order() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> list[str]: pending: list[task.Task[str]] = [ - ctx.call_activity("work", input=i, return_type=str) for i in range(2)] + ctx.as_task_context().call_activity("work", input=i, return_type=str) for i in range(2)] return await task.when_all(pending) old = _start() @@ -370,7 +392,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[str]: @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: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: try: await ctx.call_activity("boom") except task.TaskFailedError: @@ -392,9 +414,9 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: @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)) + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: + event = ctx.as_task_context().wait_for_external_event("approval") + timer = ctx.as_task_context().create_timer(timedelta(seconds=1)) winner = await task.when_any([event, timer]) if winner is event: timer.cancel() @@ -410,7 +432,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: def test_sub_orchestration_and_replay_flags() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: 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] @@ -422,7 +444,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> list[Any]: def test_retry_timer_resumes_same_coroutine() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: return await ctx.call_activity( "retry", retry_policy=task.RetryPolicy( first_retry_interval=timedelta(seconds=1), max_number_of_attempts=2)) @@ -450,7 +472,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: @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: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: retry = task.RetryPolicy(first_retry_interval=timedelta(seconds=1), max_number_of_attempts=3) try: if sub_orchestration: @@ -491,7 +513,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: def test_entity_lock_and_call_suspend_coroutine() -> None: entity = entities.EntityInstanceId("counter", "one") - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: with await ctx.lock_entities([entity]): return await ctx.call_entity(entity, "get") @@ -513,7 +535,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: def test_continue_as_new_from_coroutine() -> None: - async def orchestrator(ctx: task.OrchestrationContext, value: int) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, value: int) -> None: await ctx.wait_for_external_event("next") ctx.continue_as_new(value + 1) @@ -524,7 +546,7 @@ async def orchestrator(ctx: task.OrchestrationContext, value: int) -> None: def test_nondeterminism_still_detected_for_coroutines() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: return await ctx.call_activity("changed") complete = _completion(_execute( @@ -535,7 +557,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: @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: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: if after_activity: await ctx.call_activity("work") await asyncio.Future() diff --git a/tests/durabletask/typing/async_context.py b/tests/durabletask/typing/async_context.py new file mode 100644 index 00000000..2ab4ccec --- /dev/null +++ b/tests/durabletask/typing/async_context.py @@ -0,0 +1,58 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""The async surface returns native coroutines with inferred result types.""" + +import asyncio +from collections.abc import Coroutine, Generator +from datetime import timedelta +from typing import Any, assert_type + +from durabletask import client, task, worker +from durabletask.entities import EntityInstanceId, EntityLock + + +def increment(ctx: task.ActivityContext, value: int) -> int: + return value + 1 + + +async def child(ctx: task.AsyncOrchestrationContext, value: int) -> str: + return str(value) + + +def generator_child(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, str]: + result = yield ctx.call_activity("work", input=value) + return str(result) + + +def check_registration(hub_worker: worker.TaskHubGrpcWorker, hub_client: client.TaskHubGrpcClient) -> None: + hub_worker.add_orchestrator(child) + hub_worker.add_orchestrator(generator_child) + hub_client.schedule_new_orchestration(child, input=1) + hub_client.schedule_new_orchestration(generator_child, input=1) + + +async def check_cold_operations(ctx: task.AsyncOrchestrationContext, entity: EntityInstanceId) -> None: + ops = ctx + assert_type(ops, task.AsyncOrchestrationContext) + activity: Coroutine[Any, Any, int] = ops.call_activity(increment, input=1) + assert_type(await activity, int) + assert_type(await ops.call_activity("named", return_type=str), str) + assert_type(await ops.call_activity("named"), Any) + assert_type(await ops.call_sub_orchestrator(child, input=1), str) + assert_type(await ops.call_sub_orchestrator(generator_child, input=1), str) + assert_type(await ops.call_sub_orchestrator("child", return_type=int), int) + assert_type(await ops.call_sub_orchestrator("child"), Any) + assert_type(await ops.call_entity(entity, "get", return_type=int), int) + assert_type(await ops.call_entity(entity, "get"), Any) + assert_type(await ops.create_timer(timedelta(seconds=1)), None) + assert_type(await ops.wait_for_external_event("event", data_type=str), str) + assert_type(await ops.wait_for_external_event("event"), Any) + assert_type(await ops.lock_entities([entity]), EntityLock) + assert_type(asyncio.create_task(ops.call_activity(increment, input=1)), asyncio.Task[int]) + async with asyncio.TaskGroup() as group: + assert_type(group.create_task(ops.wait_for_external_event("event", data_type=str)), asyncio.Task[str]) + operation = ctx.as_task_context().call_activity(increment, input=1) + assert_type(operation, task.CompletableTask[int]) + assert_type(await operation.wait(), int) + assert_type(asyncio.create_task(operation.wait()), asyncio.Task[int]) From bf3cf4bf7deef238da399d398eb385664f7b6982 Mon Sep 17 00:00:00 2001 From: Andy Staples Date: Thu, 12 Nov 2026 11:32:13 -0700 Subject: [PATCH 2/5] Adapt replay regressions to cold async context Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../durabletask/test_asyncio_orchestrators.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py index a55ec596..bbf8405c 100644 --- a/tests/durabletask/test_asyncio_orchestrators.py +++ b/tests/durabletask/test_asyncio_orchestrators.py @@ -500,7 +500,7 @@ async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: def test_call_at_deadline_is_exact_microsecond_ceiling(offset: float) -> None: requested: list[float] = [] - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: loop = asyncio.get_running_loop() when = loop.time() + offset requested.append(when) @@ -523,7 +523,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: def test_call_at_known_float_deadline_rounds_up_to_next_microsecond() -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: loop = asyncio.get_running_loop() fired = loop.create_future() loop.call_at(loop.time() + 0.000003, fired.set_result, None) @@ -535,7 +535,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: @pytest.mark.parametrize("offset", [0.0, -1.0]) def test_call_at_due_deadline_runs_without_durable_timer(offset: float) -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: loop = asyncio.get_running_loop() fired = loop.create_future() loop.call_at(loop.time() + offset, fired.set_result, "fired") @@ -547,7 +547,7 @@ async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> str: @pytest.mark.parametrize("when", [1e20, float(sys.float_info.max)]) def test_call_at_out_of_range_deadline_fails(when: float) -> None: - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: asyncio.get_running_loop().call_at(when, lambda: None) complete = _completion(_execute(orchestrator, [], _start())) @@ -660,7 +660,7 @@ async def branch(index: int) -> None: async def _shared_future_cancellation_orchestrator( - ctx: task.OrchestrationContext, _: Any, *, direct_future: bool, gather_wait: bool, + ctx: task.AsyncOrchestrationContext, _: Any, *, direct_future: bool, gather_wait: bool, ) -> str: loop = asyncio.get_running_loop() futures: list[asyncio.Future[None]] = [ @@ -668,7 +668,7 @@ async def _shared_future_cancellation_orchestrator( def on_cancelled(index: int, future: asyncio.Future[None]) -> None: assert future.cancelled() - ctx.call_activity(f"cancelled-{index}") + ctx.as_task_context().call_activity(f"cancelled-{index}") for index, future in enumerate(futures): future.add_done_callback(partial(on_cancelled, index)) @@ -1083,7 +1083,7 @@ def test_cleanup_keeps_process_control_over_other_control_flow( errors.reverse() cleaned: list[int] = [] - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def branch(index: int) -> None: try: await ctx.wait_for_external_event(f"never-{index}") @@ -1106,7 +1106,7 @@ def test_cleanup_closes_async_generators_started_during_generator_cleanup( count: int, caplog: pytest.LogCaptureFixture) -> None: closed: list[int] = [] - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def resource(index: int) -> AsyncIterator[int]: try: yield 1 @@ -1148,7 +1148,7 @@ def emit(self, record): logging.getLogger().addHandler(Capture()) -async def orchestrator(ctx: task.OrchestrationContext, _): +async def orchestrator(ctx: task.AsyncOrchestrationContext, _): async def resource(): try: yield 1 @@ -1191,7 +1191,7 @@ async def signal_on_forced_close(): 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 orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: async def values() -> AsyncIterator[int]: yield 1 @@ -1217,7 +1217,7 @@ def test_entity_lock_disposal_during_teardown_does_not_schedule_releases( first = entities.EntityInstanceId("counter", "a") second = entities.EntityInstanceId("counter", "b") - async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> Any: + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: with await ctx.lock_entities([first, second]): if suspend == "activity": return await ctx.call_activity("inside-lock") From 8b769dadbe617fac6b0b70d5b9b2eaecfdae3d28 Mon Sep 17 00:00:00 2001 From: Andy Staples Date: Thu, 12 Nov 2026 13:24:53 -0700 Subject: [PATCH 3/5] Correct async sub-orchestration result typing Expose Any when no reconstruction type is supplied, preserving existing decoding behavior and explicit return_type inference. Cover typed and raw results through replay and public typing fixtures. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 4 ++ docs/async-orchestrators.md | 6 +++ durabletask/async_context.py | 9 ++++- tests/durabletask/test_async_context.py | 47 +++++++++++++++++++++++ tests/durabletask/typing/async_context.py | 34 ++++++++++++++-- 5 files changed, 95 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 62e756b1..d2c46962 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,10 @@ 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 sub-orchestration calls without an explicit `return_type` now expose +`Any` rather than promising the child function's annotated result type. +Custom result reconstruction still requires `return_type`; runtime decoding is +unchanged. - 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. diff --git a/docs/async-orchestrators.md b/docs/async-orchestrators.md index ac386350..19251f55 100644 --- a/docs/async-orchestrators.md +++ b/docs/async-orchestrators.md @@ -128,6 +128,12 @@ where supported by the provider. Functions compatibility contexts also provide async `task_all` / `task_any`. Generator and coroutine parents/children may call one another through durable sub-orchestrations. +For sub-orchestration results, pass `return_type=Result` to reconstruct a custom +result type. Unlike activity calls, a child orchestrator's return annotation is +not used for result reconstruction, even when passing the child function rather +than its name. Without an explicit `return_type`, the async API returns `Any`; +with the default JSON converter, a plain JSON object remains a dictionary. + The dispatch-local loop supports FIFO ready callbacks, local futures and tasks, `asyncio.create_task`, `asyncio.gather`, `asyncio.sleep`, `asyncio.TaskGroup`, and `asyncio.timeout` / `timeout_at`. diff --git a/durabletask/async_context.py b/durabletask/async_context.py index d8d293e7..f81d3820 100644 --- a/durabletask/async_context.py +++ b/durabletask/async_context.py @@ -157,7 +157,7 @@ async def call_sub_orchestrator( retry_policy: task.RetryPolicy | None = ..., version: str | None = ..., return_type: None = ..., - ) -> TOutput: + ) -> Any: ... @overload @@ -178,7 +178,12 @@ async def call_sub_orchestrator( version: str | None = None, return_type: type | None = None, ) -> Any: - """Start and await a sub-orchestration when this coroutine executes.""" + """Start and await a sub-orchestration when this coroutine executes. + + Supply ``return_type`` to reconstruct a custom result type. A callable's + return annotation is not used to reconstruct sub-orchestration results; + without an explicit type, the result is typed as ``Any``. + """ return await self._context.call_sub_orchestrator( orchestrator, input=input, instance_id=instance_id, retry_policy=retry_policy, version=version, return_type=return_type) diff --git a/tests/durabletask/test_async_context.py b/tests/durabletask/test_async_context.py index d822f2c2..01ff42b9 100644 --- a/tests/durabletask/test_async_context.py +++ b/tests/durabletask/test_async_context.py @@ -8,6 +8,7 @@ import json import logging from collections.abc import Coroutine, Generator +from dataclasses import dataclass from datetime import timedelta from typing import Any, get_type_hints @@ -33,9 +34,55 @@ def test_additive_public_api_has_no_new_abstract_requirements() -> None: def test_public_async_context_annotations_resolve_at_runtime() -> None: assert get_type_hints(AsyncOrchestrationContext.call_activity)["return"] is Any + assert get_type_hints(AsyncOrchestrationContext.call_sub_orchestrator)["return"] is Any assert get_type_hints(AsyncOrchestrationContext.as_task_context)["return"] is task.OrchestrationContext +@pytest.mark.parametrize("child_style", ["async", "generator", "name"]) +@pytest.mark.parametrize("reconstruction", ["omitted", "none", "explicit"]) +def test_sub_orchestration_result_requires_explicit_reconstruction_type( + child_style: str, reconstruction: str, +) -> None: + @dataclass + class Result: + value: int + + async def async_child(ctx: task.AsyncOrchestrationContext, value: int) -> Result: + return Result(value) + + def generator_child(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, Result]: + yield ctx.call_activity("work") + return Result(value) + + child = {"async": async_child, "generator": generator_child, "name": "child"}[child_style] + child_name = child if isinstance(child, str) else task.get_name(child) + + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: + if reconstruction == "explicit": + result = await ctx.call_sub_orchestrator(child, input=42, return_type=Result) + assert isinstance(result, Result) + return result.value + raw: dict[str, int] + if reconstruction == "none": + raw = await ctx.call_sub_orchestrator(child, input=42, return_type=None) + else: + raw = await ctx.call_sub_orchestrator(child, input=42) + assert isinstance(raw, dict) + return raw["value"] + + start = _start() + initial = _execute(orchestrator, [], start) + assert len(initial.actions) == 1 + scheduled = initial.actions[0] + assert scheduled.createSubOrchestration.name == child_name + history = start + [helpers.new_sub_orchestration_created_event( + scheduled.id, child_name, scheduled.createSubOrchestration.instanceId)] + complete = _completion(_execute(orchestrator, history, [ + helpers.new_sub_orchestration_completed_event(scheduled.id, '{"value": 42}')])) + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED + assert json.loads(complete.result.value) == 42 + + @pytest.mark.parametrize("cancel", [False, True]) def test_all_operations_are_cold_and_prestart_cancellation_schedules_nothing(cancel: bool) -> None: async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> Any: diff --git a/tests/durabletask/typing/async_context.py b/tests/durabletask/typing/async_context.py index 2ab4ccec..487fe27d 100644 --- a/tests/durabletask/typing/async_context.py +++ b/tests/durabletask/typing/async_context.py @@ -5,6 +5,7 @@ import asyncio from collections.abc import Coroutine, Generator +from dataclasses import dataclass from datetime import timedelta from typing import Any, assert_type @@ -20,6 +21,19 @@ async def child(ctx: task.AsyncOrchestrationContext, value: int) -> str: return str(value) +@dataclass +class Result: + value: int + + +async def result_child(ctx: task.AsyncOrchestrationContext, value: int) -> Result: + return Result(value) + + +async def none_child(ctx: task.AsyncOrchestrationContext, value: int) -> None: + pass + + def generator_child(ctx: task.OrchestrationContext, value: int) -> Generator[task.Task[Any], Any, str]: result = yield ctx.call_activity("work", input=value) return str(result) @@ -32,17 +46,31 @@ def check_registration(hub_worker: worker.TaskHubGrpcWorker, hub_client: client. hub_client.schedule_new_orchestration(generator_child, input=1) -async def check_cold_operations(ctx: task.AsyncOrchestrationContext, entity: EntityInstanceId) -> None: +async def check_cold_operations( + ctx: task.AsyncOrchestrationContext, entity: EntityInstanceId, + result_type: type[Result] | None = None, +) -> None: ops = ctx assert_type(ops, task.AsyncOrchestrationContext) activity: Coroutine[Any, Any, int] = ops.call_activity(increment, input=1) assert_type(await activity, int) assert_type(await ops.call_activity("named", return_type=str), str) assert_type(await ops.call_activity("named"), Any) - assert_type(await ops.call_sub_orchestrator(child, input=1), str) - assert_type(await ops.call_sub_orchestrator(generator_child, input=1), str) + assert_type(await ops.call_sub_orchestrator(child, input=1), Any) + assert_type(await ops.call_sub_orchestrator(child, input=1, return_type=str), str) + assert_type(await ops.call_sub_orchestrator(generator_child, input=1), Any) + assert_type(await ops.call_sub_orchestrator(generator_child, input=1, return_type=str), str) + assert_type(await ops.call_sub_orchestrator(result_child, input=1), Any) + assert_type(await ops.call_sub_orchestrator(result_child, input=1, return_type=None), Any) + assert_type(await ops.call_sub_orchestrator(result_child, input=1, return_type=Result), Result) + assert_type(await ops.call_sub_orchestrator(result_child, input=1, return_type=result_type), Result | Any) + assert_type(await ops.call_sub_orchestrator("child", return_type=Result), Result) assert_type(await ops.call_sub_orchestrator("child", return_type=int), int) assert_type(await ops.call_sub_orchestrator("child"), Any) + assert_type(await ops.call_sub_orchestrator("child", return_type=None), Any) + assert_type(await ops.call_sub_orchestrator("child", return_type=result_type), Result | Any) + assert_type(await ops.call_sub_orchestrator(none_child, input=1), Any) + assert_type(await ops.call_sub_orchestrator(none_child, input=1, return_type=type(None)), None) assert_type(await ops.call_entity(entity, "get", return_type=int), int) assert_type(await ops.call_entity(entity, "get"), Any) assert_type(await ops.create_timer(timedelta(seconds=1)), None) From f8347b8759b865a7dd3f08eecad2991d5d801526 Mon Sep 17 00:00:00 2001 From: Andy Staples Date: Thu, 12 Nov 2026 13:46:12 -0700 Subject: [PATCH 4/5] Check cold-context routing in wrapped orchestrator regressions Adapt the new parent classifier tests to the async context introduced by this layer, preserving their callable-chain coverage and asserting runtime context selection. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../test_coroutine_orchestrators.py | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/durabletask/test_coroutine_orchestrators.py b/tests/durabletask/test_coroutine_orchestrators.py index 9c2d5dbb..095bc068 100644 --- a/tests/durabletask/test_coroutine_orchestrators.py +++ b/tests/durabletask/test_coroutine_orchestrators.py @@ -103,46 +103,48 @@ async def __call__( "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: + def base(ctx: task.AsyncOrchestrationContext, value: int) -> int: + assert isinstance(ctx, task.AsyncOrchestrationContext) return value + 1 @wraps(base) - async def asynchronous(ctx: task.OrchestrationContext, value: int) -> int: + async def asynchronous(ctx: task.AsyncOrchestrationContext, value: int) -> int: return base(ctx, value) @wraps(asynchronous) - def traced(ctx: task.OrchestrationContext, value: int) -> Any: + def traced(ctx: task.AsyncOrchestrationContext, value: int) -> Any: return asynchronous(ctx, value) class DecoratedCallable: - def base(self, ctx: task.OrchestrationContext, value: int) -> int: + def base(self, ctx: task.AsyncOrchestrationContext, value: int) -> int: + assert isinstance(ctx, task.AsyncOrchestrationContext) return value + 1 - async def asynchronous(self, ctx: task.OrchestrationContext, value: int) -> int: + async def asynchronous(self, ctx: task.AsyncOrchestrationContext, value: int) -> int: return self.base(ctx, value) - def __call__(self, ctx: task.OrchestrationContext, value: int) -> Any: + def __call__(self, ctx: task.AsyncOrchestrationContext, 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: + async def __call__(self, ctx: task.AsyncOrchestrationContext, 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]: + def wrap_callable(target: task.AsyncOrchestrator[Any, Any]) -> task.AsyncOrchestrator[Any, Any]: class Wrapper: - def __call__(self, ctx: task.OrchestrationContext, value: int) -> Any: + def __call__(self, ctx: task.AsyncOrchestrationContext, value: int) -> Any: return target(ctx, value) update_wrapper(Wrapper.__call__, target) return Wrapper() decorated = DecoratedCallable() - nested: task.Orchestrator[Any, Any] = traced + nested: task.AsyncOrchestrator[Any, Any] = traced # Each __call__ lookup produces a temporary bound method. for _ in range(32): nested = wrap_callable(nested) From 3a684c6e4951d59ea6a182e27af394f1f58d4263 Mon Sep 17 00:00:00 2001 From: Andy Staples Date: Thu, 12 Nov 2026 16:16:00 -0700 Subject: [PATCH 5/5] Add async context parity checks and reject non-finite timer delays Guard AsyncOrchestrationContext against public API and metadata drift from OrchestrationContext with an explicit allowlist, reject non-finite create_timer delays like event timeouts, and cover the race-winner, completion-check, and result idioms. Document them along with un-awaited operations. Remove the changelog entry for unreleased behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 4 - docs/async-orchestrators.md | 23 ++++ durabletask/async_context.py | 2 + tests/durabletask/test_async_context.py | 108 ++++++++++++++++++ .../durabletask/test_asyncio_orchestrators.py | 2 +- 5 files changed, 134 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 88a684a2..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 sub-orchestration calls without an explicit `return_type` now expose -`Any` rather than promising the child function's annotated result type. -Custom result reconstruction still requires `return_type`; runtime decoding is -unchanged. - 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 e76975b2..8de9b64d 100644 --- a/docs/async-orchestrators.md +++ b/docs/async-orchestrators.md @@ -100,6 +100,8 @@ allocate action IDs, serialize input, or consume buffered events. Cancelling its its first step starts no durable operation. Operation inputs are not copied: mutations made before the coroutine's first step are observed, and relative timer delays start from that execution point. Keep or copy inputs accordingly. +Dropping an operation without awaiting it or wrapping it in a task schedules +nothing, and Python reports that the coroutine was never awaited. Cold coroutines are single-use, like ordinary Python coroutines. For an explicitly started operation with multiple independent waiters, retain its @@ -118,6 +120,27 @@ Cancelling either borrower does not cancel the underlying operation. Direct an awaitable durable handle is not a coroutine. Do not retry that expression assuming the eager operation was rolled back. +To identify the winner of a race, check whether an operation has completed, or +read its result later, use eager handles with `task.when_any`, or wrap cold +operations in `asyncio` tasks: + +```python +eager = ctx.as_task_context() +approval = eager.wait_for_external_event("approval", data_type=bool) +deadline = eager.create_timer(timedelta(hours=1)) +winner = await task.when_any([approval, deadline]) # The identical handle wins. +if winner is approval: + deadline.cancel() + approved = approval.get_result() +else: + report = asyncio.create_task(ctx.call_activity("report", return_type=str)) + finished = report.done() # Completion check without waiting. + summary = await report # Result, or the operation's failure. +``` + +The losing handle stays inspectable through `is_complete` and, once complete, +`get_result()`. + The worker classifies each orchestrator once, when it is registered: 1. A callable that is itself a coroutine function (including partials, bound diff --git a/durabletask/async_context.py b/durabletask/async_context.py index f81d3820..f200a6c4 100644 --- a/durabletask/async_context.py +++ b/durabletask/async_context.py @@ -230,6 +230,8 @@ async def create_timer(self, fire_at: datetime | timedelta) -> None: coroutine is constructed. Cancellation removes uncommitted timer work where possible; it cannot retract a timer already present in history. """ + if isinstance(fire_at, timedelta) and not math.isfinite(fire_at.total_seconds()): + raise ValueError("Timer delays must be finite.") timer = self._context.create_timer(fire_at) try: await timer diff --git a/tests/durabletask/test_async_context.py b/tests/durabletask/test_async_context.py index 01ff42b9..8ee75850 100644 --- a/tests/durabletask/test_async_context.py +++ b/tests/durabletask/test_async_context.py @@ -19,6 +19,7 @@ from durabletask.internal import helpers from durabletask.internal import orchestrator_service_pb2 as pb from tests.durabletask.test_coroutine_orchestrators import ( + INSTANCE_ID, NOW, _completion, # pyright: ignore[reportPrivateUsage] _execute, # pyright: ignore[reportPrivateUsage] @@ -32,6 +33,56 @@ def test_additive_public_api_has_no_new_abstract_requirements() -> None: assert task.AsyncOrchestrationContext is AsyncOrchestrationContext +# Intentional differences between the contexts; anything else is drift. +_ASYNC_ONLY_MEMBERS = {"as_task_context"} +_EAGER_ONLY_MEMBERS: set[str] = set() +_ASYNC_OPERATIONS = { + "call_activity", "call_sub_orchestrator", "call_entity", "lock_entities", + "create_timer", "wait_for_external_event", +} +_EXTRA_ASYNC_PARAMETERS = {"wait_for_external_event": {"timeout"}} + + +def _public_members(cls: type) -> set[str]: + return {name for name in dir(cls) if not name.startswith("_")} + + +def test_async_context_public_api_matches_task_context() -> None: + async_members = _public_members(AsyncOrchestrationContext) + eager_members = _public_members(task.OrchestrationContext) + assert async_members - eager_members == _ASYNC_ONLY_MEMBERS + assert eager_members - async_members == _EAGER_ONLY_MEMBERS + + for name in sorted(async_members & eager_members): + async_member = inspect.getattr_static(AsyncOrchestrationContext, name) + eager_member = inspect.getattr_static(task.OrchestrationContext, name) + assert isinstance(async_member, property) == isinstance(eager_member, property), name + if isinstance(async_member, property): + continue + assert inspect.iscoroutinefunction(async_member) == (name in _ASYNC_OPERATIONS), name + async_parameters = set(inspect.signature(async_member).parameters) + eager_parameters = set(inspect.signature(eager_member).parameters) + assert async_parameters - eager_parameters == _EXTRA_ASYNC_PARAMETERS.get(name, set()), name + assert eager_parameters <= async_parameters, name + + +def test_async_context_metadata_matches_task_context() -> None: + observed: list[tuple[Any, Any]] = [] + + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: + eager = ctx.as_task_context() + for name in ("instance_id", "parent_instance_id", "version", "current_utc_datetime", "is_replaying"): + observed.append((getattr(ctx, name), getattr(eager, name))) + await ctx.call_activity("work") + observed.append((ctx.is_replaying, eager.is_replaying)) + + old = _start() + [helpers.new_task_scheduled_event(1, "work")] + _execute(orchestrator, old, [helpers.new_task_completed_event(1, "null")]) + assert observed and all(async_value == eager_value for async_value, eager_value in observed) + assert observed[0][0] == INSTANCE_ID + assert observed[4][0] is True and observed[-1][0] is False + + def test_public_async_context_annotations_resolve_at_runtime() -> None: assert get_type_hints(AsyncOrchestrationContext.call_activity)["return"] is Any assert get_type_hints(AsyncOrchestrationContext.call_sub_orchestrator)["return"] is Any @@ -348,6 +399,63 @@ async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> int: assert json.loads(_completion(result).result.value) == 42 +@pytest.mark.parametrize("seconds", [float("inf"), float("-inf"), float("nan")]) +def test_nonfinite_timer_delay_rejected_before_scheduling(seconds: float) -> None: + class InfiniteDelay(timedelta): + def total_seconds(self) -> float: + return seconds + + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: + with pytest.raises(ValueError, match="finite"): + await ctx.create_timer(InfiniteDelay()) + return await ctx.call_activity("work", return_type=str) + + actions = _execute(orchestrator, [], _start()).actions + assert [(action.id, action.scheduleTask.name) for action in actions] == [(1, "work")] + + +@pytest.mark.parametrize("approved_first", [True, False]) +def test_winner_completion_and_result_idioms(approved_first: bool) -> None: + observed: list[Any] = [] + + async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> str: + eager = ctx.as_task_context() + approval = eager.wait_for_external_event("approval", data_type=bool) + deadline = eager.create_timer(timedelta(hours=1)) + winner = await task.when_any([approval, deadline]) + observed.append((winner is approval, approval.is_complete, deadline.is_complete)) + if winner is approval: + deadline.cancel() + return f"approved:{approval.get_result()}" + + report = asyncio.create_task(ctx.call_activity("report", return_type=str)) + await asyncio.sleep(0) + observed.append(report.done()) + return f"timed-out:{await report}" + + deadline = NOW + timedelta(hours=1) + old = _start() + [helpers.new_timer_created_event(1, deadline)] + if approved_first: + new = [helpers.new_event_raised_event("approval", "true")] + for _ in range(2): + observed.clear() + complete = _completion(_execute(orchestrator, old, new)) + assert json.loads(complete.result.value) == "approved:True" + assert observed == [(True, True, False)] + return + + fired = [helpers.new_orchestrator_started_event(deadline), helpers.new_timer_fired_event(1, deadline)] + actions = _execute(orchestrator, old, fired).actions + assert [action.scheduleTask.name for action in actions] == ["report"] + assert observed == [(False, False, True), False] + old += fired + [helpers.new_task_scheduled_event(2, "report")] + for _ in range(2): + observed.clear() + complete = _completion(_execute(orchestrator, old, [helpers.new_task_completed_event(2, '"sent"')])) + assert json.loads(complete.result.value) == "timed-out:sent" + assert observed == [(False, False, True), False] + + def test_owned_timer_cancellation_does_not_cancel_borrowed_timer() -> None: async def orchestrator(ctx: task.AsyncOrchestrationContext, _: Any) -> None: borrowed = ctx.as_task_context().create_timer(timedelta(seconds=10)) diff --git a/tests/durabletask/test_asyncio_orchestrators.py b/tests/durabletask/test_asyncio_orchestrators.py index 026544b6..d5dcb0ea 100644 --- a/tests/durabletask/test_asyncio_orchestrators.py +++ b/tests/durabletask/test_asyncio_orchestrators.py @@ -646,7 +646,7 @@ def test_untracked_task_finalized_after_dispatch_close_is_quiet(monkeypatch: pyt untracked: list[asyncio.Task[Any]] = [] async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: - untracked.append(asyncio.Task(ctx.call_activity("background").wait())) + untracked.append(asyncio.Task(ctx.call_activity("background"))) await ctx.call_activity("work") _execute(orchestrator, [], _start())