diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 0408e021ca1..719e21f0073 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -47,12 +47,13 @@ import json import logging import math +import sys import typing from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence from contextvars import Context, ContextVar from copy import deepcopy from types import CodeType -from typing import Any, Generic, Literal, TypeVar, overload +from typing import Any, Generic, Literal, TypeVar, cast, overload from .._agents import BaseAgent from .._feature_stage import ExperimentalFeature, experimental @@ -101,15 +102,33 @@ def _create_task( coro: Any, **kwargs: Any, ) -> asyncio.Future[Any]: - if self.original_factory is None: - task = asyncio.Task(coro, loop=loop, **kwargs) - else: - task = self.original_factory(loop, coro, **kwargs) - + parent_task = asyncio.current_task(loop=loop) task_context = kwargs.get("context") ctx = task_context.get(_active_run_ctx) if isinstance(task_context, Context) else _active_run_ctx.get() - if ctx is not None: + if ctx is None: + if self.original_factory is None: + return asyncio.Task(coro, loop=loop, **kwargs) + return self.original_factory(loop, coro, **kwargs) + + def _register(task: asyncio.Future[Any]) -> None: ctx._workflow_tasks.add(task) + ctx._workflow_task_parents[task] = parent_task + if parent_task in ctx._pause_owned_tasks: + ctx._pause_owned_tasks.add(task) + if ctx._workflow_interruption is not None and parent_task in ctx._active_step_tasks: + ctx._pause_step_entry_tasks.add(task) + + async def _tracked_coro() -> Any: + current_task = asyncio.current_task() + if current_task is not None: + _register(current_task) + return await coro + + if self.original_factory is None: + task = asyncio.Task(_tracked_coro(), loop=loop, **kwargs) + else: + task = self.original_factory(loop, _tracked_coro(), **kwargs) + _register(task) return task @@ -410,6 +429,105 @@ def _snapshot_replayed_step_cache_value(value: Any) -> Any: raise ValueError("@step result cannot be safely copied for replay.") +def _prepare_workflow_state_snapshot( + value: Any, + memo: dict[int, Any], + seen: dict[int, bool], + *, + hash_required: bool = False, +) -> None: + """Pre-register non-copyable leaves before standard deepcopy.""" + value_id = id(value) + if value_id in memo: + if hash_required: + try: + hash(memo[value_id]) + except Exception: + memo[value_id] = value + return + previous_requirement = seen.get(value_id) + if previous_requirement is True or (previous_requirement is False and not hash_required): + return + seen[value_id] = hash_required or bool(previous_requirement) + + if isinstance(value, dict): + for key, item in cast(dict[Any, Any], value).items(): + _prepare_workflow_state_snapshot(key, memo, seen, hash_required=True) + _prepare_workflow_state_snapshot(item, memo, seen) + if type(cast(Any, value)) is dict: + return + + if isinstance(value, list): + for item in cast(list[Any], value): + _prepare_workflow_state_snapshot(item, memo, seen) + if type(cast(Any, value)) is list: + return + + if isinstance(value, tuple): + for item in cast(tuple[Any, ...], value): + _prepare_workflow_state_snapshot(item, memo, seen, hash_required=hash_required) + if type(cast(Any, value)) is tuple: + return + if isinstance(value, set): + for item in cast(set[Any], value): + _prepare_workflow_state_snapshot(item, memo, seen, hash_required=True) + if type(cast(Any, value)) is set: + return + if isinstance(value, frozenset): + for item in cast(frozenset[Any], value): + _prepare_workflow_state_snapshot(item, memo, seen, hash_required=True) + if type(cast(Any, value)) is frozenset: + return + + value_any = cast(Any, value) + if getattr(value_any.__class__, "__deepcopy__", None) is None: + has_object_state = False + object_dict = getattr(cast(Any, value), "__dict__", None) + if isinstance(object_dict, dict): + has_object_state = True + for item in cast(dict[str, Any], object_dict).values(): + _prepare_workflow_state_snapshot(item, memo, seen) + + for cls in value_any.__class__.__mro__: + raw_slots = cls.__dict__.get("__slots__", ()) + slots = (raw_slots,) if isinstance(raw_slots, str) else raw_slots + for raw_slot in slots: + if raw_slot in ("__dict__", "__weakref__"): + continue + slot = ( + f"_{cls.__name__.lstrip('_')}{raw_slot}" + if raw_slot.startswith("__") and not raw_slot.endswith("__") + else raw_slot + ) + try: + item = getattr(cast(Any, value), slot) + except AttributeError: + continue + has_object_state = True + _prepare_workflow_state_snapshot(item, memo, seen) + + if has_object_state: + return + + candidate_memo = dict(memo) + try: + snapshot: Any = deepcopy(value_any, candidate_memo) + if hash_required: + hash(snapshot) + except Exception: + memo[value_id] = value + else: + memo.update(candidate_memo) + + +def _snapshot_workflow_state(state: Mapping[str, Any]) -> dict[str, Any]: + """Snapshot state while retaining only specific non-copyable leaves.""" + state_dict = dict(state) + memo: dict[int, Any] = {} + _prepare_workflow_state_snapshot(state_dict, memo, {}) + return deepcopy(state_dict, memo) + + @experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) def get_run_context() -> RunContext | None: """Return the active :class:`RunContext`, or ``None`` if not inside a ``@workflow``. @@ -426,12 +544,24 @@ def get_run_context() -> RunContext | None: # --------------------------------------------------------------------------- -class WorkflowInterrupted(BaseException): +if sys.version_info >= (3, 11): + + class _WorkflowInterruptedBase(asyncio.CancelledError): + pass + +else: + + class _WorkflowInterruptedBase(BaseException): + pass + + +class WorkflowInterrupted(_WorkflowInterruptedBase): """Internal: raised when request_info() is called during initial execution. - Inherits from ``BaseException`` (not ``Exception``) so that user code - with ``except Exception:`` handlers inside a ``@workflow`` function does - not accidentally intercept the HITL interruption signal. + On Python 3.11+, the signal is cancellation-shaped so ``TaskGroup`` does + not cancel sibling tasks. Python 3.10 retains the original ``BaseException`` + behavior because ``gather`` normalizes cancellation subclasses there. + User code with ``except Exception:`` handlers does not intercept either form. """ def __init__(self, request_id: str, request_data: Any, response_type: type) -> None: @@ -441,6 +571,45 @@ def __init__(self, request_id: str, request_data: Any, response_type: type) -> N super().__init__(f"Workflow interrupted by request_info (request_id={request_id})") +def _get_workflow_interruption(exc: BaseException) -> WorkflowInterrupted | None: + """Return an interruption when every leaf in *exc* is an internal interruption.""" + if isinstance(exc, WorkflowInterrupted): + return exc + if type(exc).__name__ not in ("ExceptionGroup", "BaseExceptionGroup"): + return None + + nested = getattr(exc, "exceptions", ()) + if not nested: + return None + + first: WorkflowInterrupted | None = None + for child in nested: + interruption = _get_workflow_interruption(child) + if interruption is None: + return None + if first is None: + first = interruption + return first + + +def _split_workflow_interruption( + exc: BaseException, +) -> tuple[WorkflowInterrupted | None, BaseException | None]: + """Separate internal workflow interruptions from any real failures.""" + if isinstance(exc, WorkflowInterrupted): + return exc, None + if type(exc).__name__ not in ("ExceptionGroup", "BaseExceptionGroup"): + return None, exc + + split = getattr(exc, "split", None) + if split is None: + return None, exc + interruptions, remaining = split(WorkflowInterrupted) + if interruptions is None: + return None, exc + return _get_workflow_interruption(interruptions), remaining + + # --------------------------------------------------------------------------- # RunContext # --------------------------------------------------------------------------- @@ -525,8 +694,35 @@ def __init__( self._used_explicit_step_cache_keys: set[str] = set() # Tasks created while this workflow context is active. self._workflow_tasks: set[asyncio.Future[Any]] = set() + # Tasks explicitly cancelled by framework cleanup. + self._framework_cancelled_tasks: set[asyncio.Future[Any]] = set() + # Parent task for each tracked workflow task. + self._workflow_task_parents: dict[asyncio.Future[Any], asyncio.Task[Any] | None] = {} + # Tasks that have entered at least one @step wrapper. + self._step_tasks: set[asyncio.Task[Any]] = set() + # Newly spawned pause-owned tasks allowed to enter their first @step. + self._pause_step_entry_tasks: set[asyncio.Future[Any]] = set() + # Reference counts for tasks currently executing one or more nested @steps. + self._active_step_tasks: dict[asyncio.Task[Any], int] = {} + # Task ownership tree rooted at steps active when request_info interrupted. + self._pause_owned_tasks: set[asyncio.Future[Any]] = set() + # Pause-owned tasks whose errors can be masked by the internal interruption. + self._pause_failure_boundaries: set[asyncio.Future[Any]] = set() + # Concurrency waiter responsible for each pause failure boundary. + self._pause_boundary_waiters: dict[asyncio.Future[Any], asyncio.Future[Any]] = {} + # Adapter future to task whose interruption caused its cancellation. + self._pause_adapter_sources: dict[asyncio.Future[Any], asyncio.Task[Any]] = {} + # Waiter the root workflow task was blocked on when request_info interrupted. + self._pause_root_waiter: asyncio.Future[Any] | None = None + # First internal interruption raised during this run, even if user concurrency swallows it. + self._workflow_interruption: WorkflowInterrupted | None = None + # First internal interruption raised in each task. + self._task_interruptions: dict[asyncio.Task[Any], WorkflowInterrupted] = {} # The task executing the workflow body; gather/task-group children differ from this task. self._root_task: asyncio.Task[Any] | None = None + # Exact source location suspended when request_info interrupted. + self._pause_origin_code: CodeType | None = None + self._pause_origin_lineno: int | None = None # Deterministic call counter for auto-generated request_info IDs self._auto_request_info_index: int = 0 @@ -581,6 +777,7 @@ async def request_info( WorkflowInterrupted: Raised internally on initial execution (not visible to workflow authors). """ + self._raise_if_current_task_interrupted() if request_id is None: # Deterministic id; same determinism contract as @step caching. rid = f"auto::{self._auto_request_info_index}" @@ -596,6 +793,8 @@ async def request_info( value = _coerce_request_info_response(value, response_type, rid) return value + self._state = _snapshot_workflow_state(self._state) + # No response — emit event and interrupt event = WorkflowEvent.request_info( request_id=rid, @@ -605,7 +804,18 @@ async def request_info( ) await self.add_event(event) self._pending_requests[rid] = event - raise WorkflowInterrupted(rid, request_data, response_type) + self._pause_owned_tasks.update(self._active_step_tasks) + self._pause_failure_boundaries.update(self._active_step_tasks) + self._capture_pause_waiter_tasks() + self._capture_root_pause_location() + self._expand_pause_owned_tasks() + interruption = WorkflowInterrupted(rid, request_data, response_type) + if self._workflow_interruption is None: + self._workflow_interruption = interruption + current_task = asyncio.current_task() + if current_task is not None: + self._task_interruptions[current_task] = interruption + raise interruption async def add_event(self, event: WorkflowEvent[Any]) -> None: """Add a custom event to the workflow event stream. @@ -616,6 +826,7 @@ async def add_event(self, event: WorkflowEvent[Any]) -> None: Args: event: The workflow event to append. """ + self._raise_if_current_task_interrupted() self._events.append(event) def get_state(self, key: str, default: Any = None) -> Any: @@ -631,6 +842,7 @@ def get_state(self, key: str, default: Any = None) -> Any: Returns: The stored value, or *default* if the key does not exist. """ + self._raise_if_current_task_interrupted() return self._state.get(key, default) def set_state(self, key: str, value: Any) -> None: @@ -650,6 +862,7 @@ def set_state(self, key: str, value: Any) -> None: ValueError: If *key* begins with ``_`` (reserved for framework bookkeeping). """ + self._raise_if_current_task_interrupted() if key.startswith("_"): raise ValueError( f"State key {key!r} starts with '_', which is reserved for " @@ -674,6 +887,305 @@ def is_streaming(self) -> bool: def _get_events(self) -> list[WorkflowEvent[Any]]: return list(self._events) + def _raise_if_current_task_interrupted(self, *, entering_step: bool = False) -> None: + """Stop framework work after the current task has requested a pause.""" + current_task = asyncio.current_task() + if current_task is not None and (interruption := self._task_interruptions.get(current_task)) is not None: + raise interruption + if self._workflow_interruption is not None and ( + current_task not in self._pause_owned_tasks + or (entering_step and current_task in self._step_tasks) + or (entering_step and current_task not in self._pause_step_entry_tasks) + or (not entering_step and current_task not in self._active_step_tasks) + ): + raise self._workflow_interruption + + def _expand_pause_owned_tasks(self) -> None: + """Include tracked descendants of tasks already owned by pause settlement.""" + while True: + owned_count = len(self._pause_owned_tasks) + self._pause_owned_tasks.update( + task for task, parent in self._workflow_task_parents.items() if parent in self._pause_owned_tasks + ) + if len(self._pause_owned_tasks) == owned_count: + return + + def _capture_waiter_tree(self, waiter: Any, seen: set[asyncio.Future[Any]]) -> None: + """Include gather children reachable through awaited task wrappers.""" + if not isinstance(waiter, asyncio.Future): + return + typed_waiter = cast(asyncio.Future[Any], waiter) + if typed_waiter in seen: + return + seen.add(typed_waiter) + children = cast(Iterable[Any], getattr(cast(Any, waiter), "_children", ())) + for task in children: + if isinstance(task, asyncio.Future): + child = cast(asyncio.Future[Any], task) + self._pause_owned_tasks.add(child) + self._pause_failure_boundaries.add(child) + self._pause_boundary_waiters[child] = typed_waiter + self._capture_waiter_tree(getattr(cast(Any, waiter), "_fut_waiter", None), seen) + + def _capture_waiter_children(self, owner: asyncio.Future[Any] | None) -> None: + """Include concurrency children reachable from the waiter awaited by *owner*.""" + self._capture_waiter_tree(getattr(owner, "_fut_waiter", None), set()) + + def _capture_gather_callback_children(self, task: asyncio.Task[Any] | None) -> None: + """Include gather siblings exposed by a child task's done callback.""" + callbacks = cast(list[Any], getattr(cast(Any, task), "_callbacks", ()) or ()) + for entry in callbacks: + raw_callback = cast(Any, entry[0] if isinstance(entry, tuple) else entry) + if not callable(raw_callback): + continue + callback = cast(Callable[..., Any], raw_callback) + if getattr(callback, "__qualname__", "") != "gather.._done_callback": + continue + try: + closure = inspect.getclosurevars(callback).nonlocals + except (TypeError, ValueError): + continue + outer = closure.get("outer") + children = closure.get("children", ()) + if not isinstance(outer, asyncio.Future): + continue + typed_outer = cast(asyncio.Future[Any], outer) + for child in cast(Iterable[Any], children): + if isinstance(child, asyncio.Future): + typed_child = cast(asyncio.Future[Any], child) + self._pause_owned_tasks.add(typed_child) + self._pause_failure_boundaries.add(typed_child) + self._pause_boundary_waiters[typed_child] = typed_outer + + def _capture_shield_adapters(self, task: asyncio.Task[Any] | None) -> None: + """Record shield futures linked to the currently interrupting task.""" + if task is None: + return + source: asyncio.Future[Any] = task + seen: set[int] = set() + while id(source) not in seen: + seen.add(id(source)) + callbacks = cast(list[Any], getattr(cast(Any, source), "_callbacks", ()) or ()) + outer: asyncio.Future[Any] | None = None + for entry in callbacks: + raw_callback = cast(Any, entry[0] if isinstance(entry, tuple) else entry) + if not callable(raw_callback): + continue + callback = cast(Callable[..., Any], raw_callback) + if getattr(callback, "__qualname__", "") != "shield.._inner_done_callback": + continue + try: + closure = inspect.getclosurevars(callback).nonlocals + except (TypeError, ValueError): + continue + candidate = closure.get("outer") + if isinstance(candidate, asyncio.Future): + outer = cast(asyncio.Future[Any], candidate) + break + if outer is None: + return + self._pause_owned_tasks.add(outer) + self._pause_failure_boundaries.add(outer) + self._pause_adapter_sources[outer] = task + source = outer + + def _capture_pause_waiter_tasks(self) -> None: + """Capture root and active-step concurrency children at the pause boundary.""" + root_waiter = getattr(self._root_task, "_fut_waiter", None) + if isinstance(root_waiter, asyncio.Future): + self._pause_root_waiter = cast(asyncio.Future[Any], root_waiter) + current_task = asyncio.current_task() + self._capture_shield_adapters(current_task) + self._capture_gather_callback_children(current_task) + self._capture_waiter_children(self._root_task) + for task in tuple(self._active_step_tasks): + self._capture_waiter_children(task) + + def _pause_boundary_was_handled(self, task: asyncio.Future[Any]) -> bool: + """Return whether the boundary's concurrency waiter handled child exceptions.""" + waiter = self._pause_boundary_waiters.get(task) + if waiter is None or not waiter.done() or waiter.cancelled(): + return False + return waiter.exception() is None + + def _pause_root_waiter_was_handled(self) -> bool: + waiter = self._pause_root_waiter + if waiter is None or not waiter.done() or waiter.cancelled(): + return False + return waiter.exception() is None + + def _capture_root_pause_location(self) -> None: + if asyncio.current_task() is self._root_task: + frame = inspect.currentframe() + while frame is not None: + if frame.f_code.co_filename != __file__: + self._pause_origin_code = frame.f_code + self._pause_origin_lineno = frame.f_lineno + return + frame = frame.f_back + return + + awaited: Any = self._root_task.get_coro() if self._root_task is not None else None + seen: set[int] = set() + candidate: tuple[CodeType, int] | None = None + while awaited is not None and id(awaited) not in seen: + seen.add(id(awaited)) + awaited_frame = getattr(awaited, "cr_frame", None) or getattr(awaited, "gi_frame", None) + if awaited_frame is not None and awaited_frame.f_code.co_filename != __file__: + candidate = (awaited_frame.f_code, awaited_frame.f_lineno) + awaited = getattr(awaited, "cr_await", None) or getattr(awaited, "gi_yieldfrom", None) + if candidate is not None: + self._pause_origin_code, self._pause_origin_lineno = candidate + + def _cancelled_at_root_pause_location(self, exc: BaseException) -> bool: + code = self._pause_origin_code + lineno = self._pause_origin_lineno + traceback = exc.__traceback__ + while code is not None and lineno is not None and traceback is not None: + if traceback.tb_frame.f_code is code and traceback.tb_lineno == lineno: + return True + traceback = traceback.tb_next + return False + + def _get_new_descendant_interruption( + self, + ancestor: asyncio.Task[Any] | None, + previous: Mapping[asyncio.Task[Any], WorkflowInterrupted], + ) -> WorkflowInterrupted | None: + """Return a newly raised interruption from *ancestor* or one of its descendants.""" + if ancestor is None: + return None + for task, interruption in self._task_interruptions.items(): + if previous.get(task) is interruption: + continue + current: asyncio.Future[Any] | None = task + while current is not None: + if current is ancestor: + return interruption + current = self._workflow_task_parents.get(current) + return None + + async def _cancel_workflow_tasks(self) -> None: + """Cancel and drain all tracked workflow tasks.""" + current_task = asyncio.current_task() + while pending := [task for task in self._workflow_tasks if task is not current_task and not task.done()]: + waiting_for_existing_cancellation = False + for task in pending: + if task in self._framework_cancelled_tasks: + continue + if self._has_pending_cancellation(task): + waiting_for_existing_cancellation = True + continue + self._cancel_workflow_task(task) + if waiting_for_existing_cancellation: + await asyncio.sleep(0) + continue + await asyncio.gather(*pending, return_exceptions=True) + tasks = [task for task in self._workflow_tasks if task is not current_task] + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + @staticmethod + def _has_pending_cancellation(task: asyncio.Future[Any]) -> bool: + waiter = getattr(task, "_fut_waiter", None) + return (isinstance(waiter, asyncio.Future) and waiter.cancelled()) or bool(getattr(task, "_must_cancel", False)) + + def _cancel_workflow_task(self, task: asyncio.Future[Any]) -> None: + """Record and request framework-owned task cancellation.""" + if task.done() or task in self._framework_cancelled_tasks or self._has_pending_cancellation(task): + return + self._framework_cancelled_tasks.add(task) + task.cancel() + + async def _settle_workflow_tasks( + self, + *, + propagate_failures: bool, + interruption: WorkflowInterrupted, + ) -> None: + """Wait for pause-associated steps, then cancel and drain detached tasks.""" + current_task = asyncio.current_task() + settled_tasks = self._pause_owned_tasks + inspected_tasks: set[asyncio.Future[Any]] = set() + + def _pending_owned_tasks() -> list[asyncio.Future[Any]]: + self._expand_pause_owned_tasks() + pending: list[asyncio.Future[Any]] = [] + for task in settled_tasks: + if task is current_task or task is self._root_task or task.done(): + continue + parent = self._workflow_task_parents.get(task) + if task in self._pause_failure_boundaries or (parent is not None and not parent.done()): + pending.append(task) + return pending + + def _raise_completed_failure() -> None: + for task in settled_tasks: + if task in inspected_tasks or not task.done(): + continue + if task is current_task or task is self._root_task: + inspected_tasks.add(task) + continue + if task.cancelled(): + inspected_tasks.add(task) + if task in self._framework_cancelled_tasks: + continue + adapter_source = self._pause_adapter_sources.get(task) + if ( + adapter_source is not None + and self._get_new_descendant_interruption(adapter_source, {}) is interruption + ): + continue + try: + task.result() + except WorkflowInterrupted: + continue + except asyncio.CancelledError as exc: + cancelling = getattr(task, "cancelling", None) + cancelling_count = cast(Callable[[], int], cancelling)() if callable(cancelling) else 0 + if ( + sys.version_info >= (3, 11) + and cancelling_count == 0 + and not exc.args + and self._get_new_descendant_interruption(cast(asyncio.Task[Any], task), {}) is interruption + ): + continue + raise + if task in self._pause_failure_boundaries: + if self._pause_boundary_was_handled(task): + inspected_tasks.add(task) + continue + else: + if not getattr(task, "_log_traceback", False): + inspected_tasks.add(task) + continue + parent = self._workflow_task_parents.get(task) + if parent is not None and not parent.done(): + continue + inspected_tasks.add(task) + error = task.exception() + if error is not None: + _, remaining = _split_workflow_interruption(error) + if remaining is not None: + raise remaining from None + + try: + while True: + if propagate_failures: + _raise_completed_failure() + pending = _pending_owned_tasks() + if not pending: + break + await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED) + if propagate_failures: + _raise_completed_failure() + await self._cancel_workflow_tasks() + if propagate_failures: + _raise_completed_failure() + except BaseException: + await self._cancel_workflow_tasks() + raise + def _get_legacy_step_cache_key(self, step_name: str) -> str: idx = self._step_call_counters.get(step_name, 0) self._step_call_counters[step_name] = idx + 1 @@ -999,6 +1511,11 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: if ctx is None: # Outside a workflow — pass through directly return await self._func(*args, **kwargs) + step_task = asyncio.current_task() + ctx._raise_if_current_task_interrupted(entering_step=True) + if step_task is not None: + ctx._step_tasks.add(step_task) + ctx._pause_step_entry_tasks.discard(step_task) legacy_cache_key = ctx._get_legacy_step_cache_key(self.name) cache_key = legacy_cache_key @@ -1035,47 +1552,61 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: "Use @step(replay_key=...) to identify each logical invocation." ) - # Inject RunContext if the step function declares it - call_args, call_kwargs = self._build_call_args_with_ctx(ctx, args, kwargs) - - # Defensive deepcopy for the event log only; fall back to the live - # reference so non-deepcopyable args (locks, sockets) don't fail. - if args or kwargs: - try: - invocation_data: Any = deepcopy({"args": args, "kwargs": kwargs}) - except Exception: - invocation_data = {"args": args, "kwargs": kwargs} - else: - invocation_data = None - await ctx.add_event(WorkflowEvent.executor_invoked(self.name, invocation_data)) - auto_request_info_index_before = ctx._auto_request_info_index + task_interruptions_before = dict(ctx._task_interruptions) + if step_task is not None: + ctx._active_step_tasks[step_task] = ctx._active_step_tasks.get(step_task, 0) + 1 try: - result = await self._func(*call_args, **call_kwargs) - except Exception as exc: - # NOTE: WorkflowInterrupted (from request_info inside a step) inherits - # from BaseException, NOT Exception, so it propagates past this handler - # without emitting a spurious executor_failed event. This is intentional - # — request_info is fully supported inside @step functions. - from ..observability import OBSERVABILITY_SETTINGS + # Inject RunContext if the step function declares it + call_args, call_kwargs = self._build_call_args_with_ctx(ctx, args, kwargs) + + # Defensive deepcopy for the event log only; fall back to the live + # reference so non-deepcopyable args (locks, sockets) don't fail. + if args or kwargs: + try: + invocation_data: Any = deepcopy({"args": args, "kwargs": kwargs}) + except Exception: + invocation_data = {"args": args, "kwargs": kwargs} + else: + invocation_data = None + await ctx.add_event(WorkflowEvent.executor_invoked(self.name, invocation_data)) + auto_request_info_index_before = ctx._auto_request_info_index + try: + result = await self._func(*call_args, **call_kwargs) + except Exception as exc: + # NOTE: WorkflowInterrupted (from request_info inside a step) inherits + # from BaseException, NOT Exception, so it propagates past this handler + # without emitting a spurious executor_failed event. This is intentional + # — request_info is fully supported inside @step functions. + from ..observability import OBSERVABILITY_SETTINGS - details = WorkflowErrorDetails.from_exception( - exc, - include_sensitive_data=OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED, + details = WorkflowErrorDetails.from_exception( + exc, + include_sensitive_data=OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED, + ) + await ctx.add_event(WorkflowEvent.executor_failed(self.name, details)) + raise + task_interruption = ctx._get_new_descendant_interruption(step_task, task_interruptions_before) + if task_interruption is not None: + raise task_interruption + ctx._set_cached_step_auto_request_info_count( + cache_key, + ctx._auto_request_info_index - auto_request_info_index_before, ) - await ctx.add_event(WorkflowEvent.executor_failed(self.name, details)) - raise - ctx._set_cached_step_auto_request_info_count( - cache_key, - ctx._auto_request_info_index - auto_request_info_index_before, - ) - copied, cached_result = _try_snapshot_step_cache_value(result) - identity_is_replayable = self._replay_key is not None or self._wrapper_identity_is_durable - ctx._set_cached_result(cache_key, cached_result, replayable=copied and identity_is_replayable) - _, event_data = _try_snapshot_step_cache_value(result) - await ctx.add_event(WorkflowEvent.executor_completed(self.name, event_data)) - if ctx._on_step_completed is not None: - await ctx._on_step_completed() - return result + copied, cached_result = _try_snapshot_step_cache_value(result) + identity_is_replayable = self._replay_key is not None or self._wrapper_identity_is_durable + ctx._set_cached_result(cache_key, cached_result, replayable=copied and identity_is_replayable) + _, event_data = _try_snapshot_step_cache_value(result) + await ctx.add_event(WorkflowEvent.executor_completed(self.name, event_data)) + if ctx._on_step_completed is not None: + await ctx._on_step_completed() + return result + finally: + if step_task is not None: + remaining_depth = ctx._active_step_tasks[step_task] - 1 + if remaining_depth: + ctx._active_step_tasks[step_task] = remaining_depth + else: + del ctx._active_step_tasks[step_task] # --------------------------------------------------------------------------- @@ -1740,7 +2271,6 @@ async def _execute(self, ctx: RunContext, message: Any) -> Any: ) token = _active_run_ctx.set(ctx) - ctx._root_task = asyncio.current_task() release_task_tracking = _track_workflow_tasks() try: sig = inspect.signature(self._func) @@ -1775,7 +2305,62 @@ async def _execute(self, ctx: RunContext, message: Any) -> Any: call_args.append(message) message_injected = True - return await self._func(*call_args) + root_task_ready = asyncio.Event() + + async def _invoke_workflow() -> Any: + await root_task_ready.wait() + return await self._func(*call_args) + + workflow_task = asyncio.create_task(_invoke_workflow()) + ctx._root_task = workflow_task + root_task_ready.set() + try: + await asyncio.wait({workflow_task}) + except BaseException: + ctx._cancel_workflow_task(workflow_task) + await asyncio.gather(workflow_task, return_exceptions=True) + await ctx._cancel_workflow_tasks() + raise + + try: + result = workflow_task.result() + except BaseException as exc: + interruption, remaining = _split_workflow_interruption(exc) + if ( + interruption is None + and isinstance(exc, asyncio.CancelledError) + and not exc.args + and ctx._workflow_interruption is not None + and not ctx._pause_root_waiter_was_handled() + and ctx._cancelled_at_root_pause_location(exc) + and sys.version_info >= (3, 11) + and workflow_task.cancelling() == 0 + ): + await ctx._settle_workflow_tasks( + propagate_failures=True, + interruption=ctx._workflow_interruption, + ) + raise ctx._workflow_interruption from None + if interruption is None: + await ctx._cancel_workflow_tasks() + raise + if remaining is not None: + await ctx._cancel_workflow_tasks() + raise remaining from None + await ctx._settle_workflow_tasks( + propagate_failures=True, + interruption=interruption, + ) + if isinstance(exc, WorkflowInterrupted): + raise + raise interruption from None + if ctx._workflow_interruption is not None: + await ctx._settle_workflow_tasks( + propagate_failures=True, + interruption=ctx._workflow_interruption, + ) + raise ctx._workflow_interruption + return result finally: release_task_tracking() _active_run_ctx.reset(token) diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 1d12399a7e4..59e4019b5df 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -9,7 +9,7 @@ import json import logging from collections.abc import Awaitable, Callable, Generator -from contextlib import contextmanager +from contextlib import contextmanager, suppress from dataclasses import dataclass from types import FunctionType from typing import Any, cast, overload @@ -1864,6 +1864,77 @@ async def wf(value: int) -> int: finally: loop.set_task_factory(original_factory) + @pytest.mark.skipif(not hasattr(asyncio, "eager_task_factory"), reason="Eager task factory is Python >= 3.12") + async def test_eager_task_factory_preserves_concurrent_replay_guard(self): + import threading + + loop = asyncio.get_running_loop() + original_factory = loop.get_task_factory() + calls = 0 + + @step + async def use_lock(lock: threading.Lock) -> str: + nonlocal calls + calls += 1 + return "done" + + @built_workflow + async def wf(_: str) -> str: + return (await asyncio.gather(use_lock(threading.Lock())))[0] + + loop.set_task_factory(cast(Any, asyncio).eager_task_factory) + try: + with pytest.raises(ValueError, match="replay_key"): + await wf.run("input") + finally: + loop.set_task_factory(original_factory) + + assert calls == 0 + + @pytest.mark.skipif(not hasattr(asyncio, "eager_task_factory"), reason="Eager task factory is Python >= 3.12") + async def test_eager_task_factory_registers_pause_owned_child_before_step_entry(self): + loop = asyncio.get_running_loop() + original_factory = loop.get_task_factory() + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + child_started = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("continue?", int, request_id="r") + + @step(replay_key=lambda: "child") + async def child() -> None: + child_started.set() + raise RuntimeError("boom") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release_sibling.wait() + await asyncio.create_task(child()) + return "unreachable" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling())) + + loop.set_task_factory(cast(Any, asyncio).eager_task_factory) + try: + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + release_sibling.set() + + with pytest.raises(RuntimeError, match="boom"): + await asyncio.wait_for(run_task, timeout=5) + finally: + loop.set_task_factory(original_factory) + + assert child_started.is_set() + async def test_versioned_identity_precedes_unrelated_legacy_entry(self): from agent_framework import WorkflowCheckpoint @@ -2611,6 +2682,1419 @@ def emit(self, record: logging.LogRecord) -> None: class TestHITLInStepWithCaching: """Regression tests: request_info inside @step combined with caching and bypass.""" + async def test_in_flight_gather_sibling_is_cached_before_pause_snapshot(self): + storage = InMemoryCheckpointStorage() + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + sibling_calls = 0 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "model") + async def model() -> str: + nonlocal sibling_calls + sibling_calls += 1 + sibling_started.set() + await release_sibling.wait() + return f"answer #{sibling_calls}" + + @built_workflow(checkpoint_storage=storage) + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), model())) + + initial_task = asyncio.ensure_future(wf.run(1)) + try: + await interrupting.wait() + await asyncio.sleep(0) + assert not initial_task.done() + finally: + release_sibling.set() + + paused = await initial_task + checkpoints = await storage.list_checkpoints(workflow_name="wf") + pause_checkpoint = checkpoints[-1] + resumed = await wf.run(responses={"y": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert "answer #1" in pause_checkpoint.state["_step_cache"].values() + assert resumed.get_outputs() == [[5, "answer #1"]] + assert sibling_calls == 1 + + async def test_cancelling_pause_settlement_cancels_tracked_sibling(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + sibling_finished = asyncio.Event() + sibling_task: asyncio.Task[Any] | None = None + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "model") + async def model() -> str: + nonlocal sibling_task + sibling_task = asyncio.current_task() + sibling_started.set() + try: + await asyncio.Event().wait() + finally: + sibling_finished.set() + return "unreachable" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), model())) + + run_task = asyncio.ensure_future(wf.run(1)) + try: + await asyncio.wait_for(interrupting.wait(), timeout=5) + await asyncio.sleep(0) + run_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + + await asyncio.wait_for(sibling_finished.wait(), timeout=5) + assert sibling_task is not None + assert sibling_task.cancelled() + finally: + if sibling_task is not None and not sibling_task.done(): + sibling_task.cancel() + await asyncio.gather(sibling_task, return_exceptions=True) + + async def test_unawaited_background_task_is_cancelled_before_pause_returns(self): + background_started = asyncio.Event() + background_finished = asyncio.Event() + background_task: asyncio.Task[None] | None = None + + async def background() -> None: + background_started.set() + try: + await asyncio.Event().wait() + finally: + background_finished.set() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> int: + nonlocal background_task + background_task = asyncio.create_task(background()) + await background_started.wait() + return await ctx.request_info("y?", int, request_id="y") + + try: + paused = await asyncio.wait_for(wf.run(1), timeout=1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert background_task is not None + assert background_task.cancelled() + assert background_finished.is_set() + finally: + if background_task is not None and not background_task.done(): + background_task.cancel() + await asyncio.gather(background_task, return_exceptions=True) + + async def test_detached_step_does_not_run_after_sibling_interrupts(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release = asyncio.Event() + late_finished = asyncio.Event() + late_calls = 0 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release.wait() + return "done" + + @step(replay_key=lambda: "late") + async def late() -> str: + nonlocal late_calls + late_calls += 1 + late_finished.set() + return f"late #{late_calls}" + + async def detached() -> None: + await release.wait() + await late() + + @built_workflow + async def wf(_: int) -> list[object]: + asyncio.create_task(detached()) + return list(await asyncio.gather(ask(), sibling())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + await asyncio.sleep(0) + release.set() + paused = await asyncio.wait_for(run_task, timeout=5) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert late_calls == 0 + + resumed = await asyncio.wait_for(wf.run(responses={"y": 5}), timeout=5) + await asyncio.wait_for(late_finished.wait(), timeout=5) + + assert resumed.get_outputs() == [[5, "done"]] + assert late_calls == 1 + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_detached_task_cannot_mutate_state_after_swallowed_interruption(self): + background_started = asyncio.Event() + release_background = asyncio.Event() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> str: + async def background() -> None: + background_started.set() + await release_background.wait() + ctx.set_state("leaked", "yes") + + asyncio.create_task(background()) + await background_started.wait() + with suppress(asyncio.CancelledError): + await ctx.request_info("continue?", int, request_id="r") + release_background.set() + await asyncio.sleep(0) + return "ignored" + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert wf._last_state == {} # pyright: ignore[reportPrivateUsage] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_preheld_mutable_state_cannot_change_after_pause(self): + background_started = asyncio.Event() + release_background = asyncio.Event() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> list[str]: + if not ctx.get_state("background_initialized", False): + ctx.set_state("background_initialized", True) + ctx.set_state("items", []) + items = ctx.get_state("items") + + async def background() -> None: + background_started.set() + await release_background.wait() + items.append("leaked") + + asyncio.create_task(background()) + await background_started.wait() + + with suppress(asyncio.CancelledError): + await ctx.request_info("continue?", int, request_id="r") + release_background.set() + await asyncio.sleep(0) + return ctx.get_state("items") + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert wf._last_state == { # pyright: ignore[reportPrivateUsage] + "background_initialized": True, + "items": [], + } + + resumed = await wf.run(responses={"r": 1}) + assert resumed.get_outputs() == [[]] + + async def test_non_deepcopyable_state_pauses_without_checkpoint_storage(self): + import threading + + lock = threading.Lock() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> str: + if ctx.get_state("lock") is None: + ctx.set_state("lock", lock) + answer = await ctx.request_info("continue?", int, request_id="r") + return f"{answer}:{ctx.get_state('lock') is lock}" + + paused = await wf.run(1) + resumed = await wf.run(responses={"r": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert resumed.get_outputs() == ["5:True"] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_nested_mutable_state_isolated_with_non_deepcopyable_leaf(self): + import threading + + background_started = asyncio.Event() + release_background = asyncio.Event() + lock = threading.Lock() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> tuple[list[str], bool]: + if ctx.get_state("bundle") is None: + items: list[str] = [] + bundle = {"items": items, "lock": lock} + ctx.set_state("bundle", bundle) + + async def background() -> None: + background_started.set() + await release_background.wait() + items.append("leaked") + + asyncio.create_task(background()) + await background_started.wait() + + with suppress(asyncio.CancelledError): + await ctx.request_info("continue?", int, request_id="r") + release_background.set() + await asyncio.sleep(0) + restored = ctx.get_state("bundle") + return restored["items"], restored["lock"] is lock + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert wf._last_state["bundle"]["items"] == [] # pyright: ignore[reportPrivateUsage] + assert wf._last_state["bundle"]["lock"] is lock # pyright: ignore[reportPrivateUsage] + + resumed = await wf.run(responses={"r": 1}) + assert resumed.get_outputs() == [([], True)] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_mutable_dict_key_isolated_with_non_deepcopyable_value(self): + import threading + + class MutableKey: + def __init__(self) -> None: + self.items: list[str] = [] + + background_started = asyncio.Event() + release_background = asyncio.Event() + key = MutableKey() + lock = threading.Lock() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> tuple[list[str], bool]: + if ctx.get_state("bundle") is None: + ctx.set_state("bundle", {key: "value", "lock": lock}) + + async def background() -> None: + background_started.set() + await release_background.wait() + key.items.append("leaked") + + asyncio.create_task(background()) + await background_started.wait() + + with suppress(asyncio.CancelledError): + await ctx.request_info("continue?", int, request_id="r") + release_background.set() + await asyncio.sleep(0) + restored = ctx.get_state("bundle") + restored_key = next(item for item in restored if isinstance(item, MutableKey)) + return restored_key.items, restored["lock"] is lock + + paused = await wf.run(1) + restored_key = next( + item + for item in wf._last_state["bundle"] # pyright: ignore[reportPrivateUsage] + if isinstance(item, MutableKey) + ) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert restored_key is not key + assert restored_key.items == [] + + resumed = await wf.run(responses={"r": 1}) + assert resumed.get_outputs() == [([], True)] + + async def test_repeated_key_with_unhashable_copy_preserves_original(self): + class HashableBadCopy: + def __deepcopy__(self, memo: dict[int, Any]) -> list[Any]: + return [] + + key = HashableBadCopy() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> bool: + if ctx.get_state("bundle") is None: + ctx.set_state("bundle", {"items": [key], "mapping": {key: "value"}}) + answer = await ctx.request_info("continue?", int, request_id="r") + assert answer == 1 + restored = ctx.get_state("bundle") + restored_key = next(iter(restored["mapping"])) + return restored["items"][0] is restored_key is key + + paused = await wf.run(1) + resumed = await wf.run(responses={"r": 1}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert resumed.get_outputs() == [True] + + async def test_tuple_alias_uses_finalized_hashable_key_snapshot(self): + class HashableBadCopy: + def __deepcopy__(self, memo: dict[int, Any]) -> list[Any]: + return [] + + key = HashableBadCopy() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> bool: + if ctx.get_state("bundle") is None: + ctx.set_state("bundle", {"items": [(key,)], "mapping": {key: "value"}}) + answer = await ctx.request_info("continue?", int, request_id="r") + assert answer == 1 + restored = ctx.get_state("bundle") + restored_key = next(iter(restored["mapping"])) + return restored["items"][0][0] is restored_key is key + + paused = await wf.run(1) + resumed = await wf.run(responses={"r": 1}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert resumed.get_outputs() == [True] + + async def test_tuple_alias_revisited_when_later_used_as_state_dict_key(self): + class HashableBadCopy: + def __deepcopy__(self, memo: dict[int, Any]) -> list[Any]: + return [] + + key = HashableBadCopy() + tuple_key = (key,) + + @built_workflow + async def wf(_: int, ctx: RunContext) -> bool: + if ctx.get_state("items") is None: + ctx.set_state("items", [tuple_key]) + ctx.set_state("mapping", {tuple_key: "value"}) + answer = await ctx.request_info("continue?", int, request_id="r") + assert answer == 1 + restored_items = ctx.get_state("items") + restored_key = next(iter(ctx.get_state("mapping"))) + return restored_items[0] is restored_key and restored_key[0] is key + + paused = await wf.run(1) + resumed = await wf.run(responses={"r": 1}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert resumed.get_outputs() == [True] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_mutable_container_subclass_isolated_with_non_deepcopyable_leaf(self): + import threading + + class Bundle(dict[str, Any]): + lock: Any + + background_started = asyncio.Event() + release_background = asyncio.Event() + lock = threading.Lock() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> tuple[list[str], bool]: + if ctx.get_state("bundle") is None: + items: list[str] = [] + bundle = Bundle(items=items) + bundle.lock = lock + ctx.set_state("bundle", bundle) + + async def background() -> None: + background_started.set() + await release_background.wait() + items.append("leaked") + + asyncio.create_task(background()) + await background_started.wait() + + with suppress(asyncio.CancelledError): + await ctx.request_info("continue?", int, request_id="r") + release_background.set() + await asyncio.sleep(0) + restored = ctx.get_state("bundle") + return restored["items"], restored.lock is lock + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert wf._last_state["bundle"]["items"] == [] # pyright: ignore[reportPrivateUsage] + + resumed = await wf.run(responses={"r": 1}) + assert resumed.get_outputs() == [([], True)] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_custom_object_state_isolated_with_non_deepcopyable_leaf(self): + import threading + + @dataclass + class Bundle: + items: list[str] + lock: Any + + background_started = asyncio.Event() + release_background = asyncio.Event() + lock = threading.Lock() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> tuple[list[str], bool]: + if ctx.get_state("bundle") is None: + items: list[str] = [] + bundle = Bundle(items=items, lock=lock) + ctx.set_state("bundle", bundle) + + async def background() -> None: + background_started.set() + await release_background.wait() + items.append("leaked") + + asyncio.create_task(background()) + await background_started.wait() + + with suppress(asyncio.CancelledError): + await ctx.request_info("continue?", int, request_id="r") + release_background.set() + await asyncio.sleep(0) + restored = ctx.get_state("bundle") + return restored.items, restored.lock is lock + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert wf._last_state["bundle"].items == [] # pyright: ignore[reportPrivateUsage] + + resumed = await wf.run(responses={"r": 1}) + assert resumed.get_outputs() == [([], True)] + + async def test_tuple_cycle_preserved_with_non_deepcopyable_leaf(self): + import threading + + lock = threading.Lock() + + @built_workflow + async def wf(_: int, ctx: RunContext) -> tuple[bool, bool]: + if ctx.get_state("cycle") is None: + cyclic_list: list[Any] = [] + cyclic_tuple = (cyclic_list, lock) + cyclic_list.append(cyclic_tuple) + ctx.set_state("cycle", cyclic_tuple) + answer = await ctx.request_info("continue?", int, request_id="r") + assert answer == 1 + restored = ctx.get_state("cycle") + return restored is restored[0][0], restored[1] is lock + + paused = await wf.run(1) + restored = wf._last_state["cycle"] # pyright: ignore[reportPrivateUsage] + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert restored is restored[0][0] + assert restored[1] is lock + + resumed = await wf.run(responses={"r": 1}) + assert resumed.get_outputs() == [(True, True)] + + async def test_pause_owned_task_cannot_mutate_state_after_step_finishes(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("continue?", int, request_id="r") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release_sibling.wait() + return "done" + + async def sibling_branch(ctx: RunContext) -> str: + result = await sibling() + ctx.set_state("leaked", "yes") + return result + + @built_workflow + async def wf(_: int, ctx: RunContext) -> list[object]: + return list(await asyncio.gather(ask(), sibling_branch(ctx))) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + release_sibling.set() + paused = await run_task + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert wf._last_state == {} # pyright: ignore[reportPrivateUsage] + + async def test_pause_owned_task_cannot_start_step_after_step_finishes(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + tail_finished = asyncio.Event() + tail_calls = 0 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("continue?", int, request_id="r") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release_sibling.wait() + return "done" + + @step(replay_key=lambda: "tail") + async def tail() -> str: + nonlocal tail_calls + tail_calls += 1 + tail_finished.set() + return f"tail #{tail_calls}" + + async def sibling_branch() -> list[str]: + return [await sibling(), await tail()] + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling_branch())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + release_sibling.set() + paused = await run_task + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert tail_calls == 0 + + resumed = await wf.run(responses={"r": 1}) + await tail_finished.wait() + + assert resumed.get_outputs() == [[1, ["done", "tail #1"]]] + assert tail_calls == 1 + + async def test_existing_pause_owned_branch_cannot_enter_first_step(self): + branch_ready = asyncio.Event() + release_branch = asyncio.Event() + tail_finished = asyncio.Event() + tail_calls = 0 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await branch_ready.wait() + return await ctx.request_info("continue?", int, request_id="r") + + @step(replay_key=lambda: "tail") + async def tail() -> str: + nonlocal tail_calls + tail_calls += 1 + tail_finished.set() + return f"tail #{tail_calls}" + + async def branch() -> str: + branch_ready.set() + await release_branch.wait() + return await tail() + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), branch())) + + run_task = asyncio.ensure_future(wf.run(1)) + await branch_ready.wait() + await asyncio.sleep(0) + release_branch.set() + paused = await asyncio.wait_for(run_task, timeout=5) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert tail_calls == 0 + + resumed = await wf.run(responses={"r": 1}) + await tail_finished.wait() + + assert resumed.get_outputs() == [[1, "tail #1"]] + assert tail_calls == 1 + + async def test_handled_task_failure_does_not_replace_pause(self): + async def fail() -> None: + raise ValueError("handled") + + @built_workflow + async def wf(_: int, ctx: RunContext) -> int: + outcome = (await asyncio.gather(fail(), return_exceptions=True))[0] + assert isinstance(outcome, ValueError) + return await ctx.request_info("y?", int, request_id="y") + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + + async def test_concurrent_handled_step_failure_does_not_replace_pause(self): + ask_started = asyncio.Event() + fail_started = asyncio.Event() + release = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + ask_started.set() + await fail_started.wait() + await release.wait() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "fail") + async def fail() -> None: + fail_started.set() + await ask_started.wait() + await release.wait() + raise RuntimeError("handled") + + @built_workflow + async def wf(_: int) -> list[object]: + tasks = asyncio.gather(ask(), fail(), return_exceptions=True) + await ask_started.wait() + await fail_started.wait() + release.set() + return list(await tasks) + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert paused.get_outputs() == [] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_swallowed_pause_still_surfaces_unhandled_sibling_failure(self): + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + async def fail() -> None: + await asyncio.sleep(0) + raise RuntimeError("boom") + + @built_workflow + async def wf(_: int) -> list[object]: + try: + return list(await asyncio.gather(ask(), fail())) + except asyncio.CancelledError: + return ["ignored"] + + with pytest.raises(RuntimeError, match="boom"): + await wf.run(1) + + async def test_nested_unhandled_sibling_failure_replaces_pause(self): + async def child_ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + async def child_fail() -> None: + await asyncio.sleep(0) + raise RuntimeError("boom") + + @step(replay_key=lambda: "outer") + async def outer() -> list[object]: + ctx = get_run_context() + assert ctx is not None + return list(await asyncio.gather(child_ask(ctx), child_fail())) + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(outer())) + + with pytest.raises(RuntimeError, match="boom"): + await wf.run(1) + + async def test_wait_for_wrapped_sibling_failure_replaces_pause(self): + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + async def fail() -> None: + await asyncio.sleep(0) + raise RuntimeError("boom") + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.wait_for(asyncio.gather(ask(), fail()), timeout=10)) + + with pytest.raises(RuntimeError, match="boom"): + await wf.run(1) + + @pytest.mark.parametrize("wrapper", ["wait_for", "shield"]) + async def test_wrapped_request_info_pauses_and_resumes(self, wrapper: str): + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + @built_workflow + async def wf(_: int) -> int: + if wrapper == "wait_for": + return await asyncio.wait_for(ask(), timeout=10) + return await asyncio.shield(ask()) + + paused = await wf.run(1) + resumed = await wf.run(responses={"y": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert resumed.get_outputs() == [5] + + async def test_gather_shield_wrapped_request_info_pauses_and_resumes(self): + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(asyncio.shield(ask()))) + + paused = await wf.run(1) + resumed = await wf.run(responses={"y": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert resumed.get_outputs() == [[5]] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_user_cancellation_after_swallowed_pause_is_not_reclassified(self): + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + @built_workflow + async def wf(_: int) -> None: + await asyncio.gather(ask(), return_exceptions=True) + raise asyncio.CancelledError("user cancel") + + with pytest.raises(asyncio.CancelledError): + await wf.run(1) + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + @pytest.mark.parametrize("mode", ["raise", "cancel"]) + async def test_empty_root_cancellation_after_swallowed_pause_is_not_reclassified(self, mode: str): + @built_workflow + async def wf(_: int, ctx: RunContext) -> None: + with suppress(asyncio.CancelledError): + await ctx.request_info("y?", int, request_id="y") + if mode == "cancel": + task = asyncio.current_task() + assert task is not None + task.cancel() + await asyncio.sleep(0) + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await wf.run(1) + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_helper_empty_cancellation_after_swallowed_pause_is_not_reclassified(self): + @built_workflow + async def wf(_: int, ctx: RunContext) -> None: + async def helper() -> None: + with suppress(asyncio.CancelledError): + await ctx.request_info("y?", int, request_id="y") + raise asyncio.CancelledError + + await helper() + + with pytest.raises(asyncio.CancelledError): + await wf.run(1) + + async def test_concurrent_unhandled_step_failure_replaces_pause(self): + ask_started = asyncio.Event() + fail_started = asyncio.Event() + release = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + ask_started.set() + await fail_started.wait() + await release.wait() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "fail") + async def fail() -> None: + fail_started.set() + await ask_started.wait() + await release.wait() + await asyncio.sleep(0) + raise RuntimeError("boom") + + @built_workflow + async def wf(_: int) -> list[object]: + tasks = asyncio.gather(ask(), fail()) + await ask_started.wait() + await fail_started.wait() + release.set() + return list(await tasks) + + with pytest.raises(RuntimeError, match="boom"): + await wf.run(1) + + async def test_concurrent_unhandled_plain_failure_replaces_pause(self): + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + async def fail() -> None: + await asyncio.sleep(0) + raise RuntimeError("boom") + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), fail())) + + with pytest.raises(RuntimeError, match="boom"): + await wf.run(1) + + async def test_concurrent_user_cancellation_replaces_pause(self): + cancel_started = asyncio.Event() + interrupting = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await cancel_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "cancel") + async def cancel() -> None: + cancel_started.set() + await interrupting.wait() + await asyncio.sleep(0) + raise asyncio.CancelledError("user cancel") + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), cancel())) + + with pytest.raises(asyncio.CancelledError): + await wf.run(1) + + async def test_explicit_sibling_task_cancel_replaces_pause(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + sibling_task: asyncio.Task[str] | None = None + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + nonlocal sibling_task + sibling_task = asyncio.current_task() + sibling_started.set() + await asyncio.Event().wait() + return "unreachable" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + assert sibling_task is not None + sibling_task.cancel("user cancel") + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + + async def test_preexisting_child_cancel_is_not_claimed_by_cleanup(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + background_started = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + async def background() -> None: + background_started.set() + await asyncio.Event().wait() + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release_sibling.wait() + background_task = asyncio.create_task(background()) + await background_started.wait() + background_task.cancel("user cancel") + return "done" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + release_sibling.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + + async def test_swallowed_child_cancel_can_be_cancelled_by_cleanup(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + background_started = asyncio.Event() + background_finished = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + async def background() -> None: + background_started.set() + with suppress(asyncio.CancelledError): + await asyncio.Event().wait() + try: + await asyncio.Event().wait() + finally: + background_finished.set() + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release_sibling.wait() + background_task = asyncio.create_task(background()) + await background_started.wait() + background_task.cancel("user cancel") + return "done" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + release_sibling.set() + paused = await asyncio.wait_for(run_task, timeout=5) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert background_finished.is_set() + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_later_task_cancel_after_swallowed_pause_is_not_suppressed(self): + second_started = asyncio.Event() + first_swallowed = asyncio.Event() + second_interrupting = asyncio.Event() + first_task: asyncio.Task[str] | None = None + + @step(replay_key=lambda: "first") + async def first(ctx: RunContext) -> str: + nonlocal first_task + first_task = asyncio.current_task() + await second_started.wait() + with suppress(asyncio.CancelledError): + await ctx.request_info("first?", int, request_id="r1") + first_swallowed.set() + await asyncio.Event().wait() + return "unreachable" + + @step(replay_key=lambda: "second") + async def second(ctx: RunContext) -> int: + second_started.set() + await first_swallowed.wait() + second_interrupting.set() + return await ctx.request_info("second?", int, request_id="r2") + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(first(), second())) + + run_task = asyncio.ensure_future(wf.run(1)) + await second_interrupting.wait() + assert first_task is not None + first_task.cancel("user cancel") + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_later_empty_cancellation_after_swallowed_pause_is_not_suppressed(self): + second_started = asyncio.Event() + first_swallowed = asyncio.Event() + second_interrupting = asyncio.Event() + release_first = asyncio.Event() + + @step(replay_key=lambda: "first") + async def first(ctx: RunContext) -> str: + await second_started.wait() + with suppress(asyncio.CancelledError): + await ctx.request_info("first?", int, request_id="r1") + first_swallowed.set() + await release_first.wait() + raise asyncio.CancelledError + + @step(replay_key=lambda: "second") + async def second(ctx: RunContext) -> int: + second_started.set() + await first_swallowed.wait() + second_interrupting.set() + return await ctx.request_info("second?", int, request_id="r2") + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(first(), second())) + + run_task = asyncio.ensure_future(wf.run(1)) + await second_interrupting.wait() + await asyncio.sleep(0) + release_first.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + + async def test_concurrent_plain_failure_cancels_hanging_sibling(self): + hanging_finished = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + async def fail() -> None: + await asyncio.sleep(0) + raise RuntimeError("boom") + + async def hang() -> None: + try: + await asyncio.Event().wait() + finally: + hanging_finished.set() + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), fail(), hang())) + + with pytest.raises(RuntimeError, match="boom"): + await asyncio.wait_for(wf.run(1), timeout=5) + + assert hanging_finished.is_set() + + async def test_step_failure_spawned_during_settlement_replaces_pause(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + child_started = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "child") + async def child() -> None: + child_started.set() + raise RuntimeError("boom") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release_sibling.wait() + asyncio.create_task(child()) + await child_started.wait() + return "done" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + await asyncio.sleep(0) + release_sibling.set() + + with pytest.raises(RuntimeError, match="boom"): + await run_task + + async def test_handled_child_failure_during_settlement_does_not_replace_pause(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + async def child() -> None: + raise RuntimeError("handled child") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + sibling_started.set() + await release_sibling.wait() + child_task = asyncio.create_task(child()) + with suppress(RuntimeError): + await child_task + return "done" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + release_sibling.set() + paused = await run_task + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + + async def test_unawaited_child_of_settled_sibling_is_cancelled(self): + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + background_started = asyncio.Event() + background_finished = asyncio.Event() + background_task: asyncio.Task[None] | None = None + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + async def background() -> None: + background_started.set() + try: + await asyncio.Event().wait() + finally: + background_finished.set() + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + nonlocal background_task + sibling_started.set() + await release_sibling.wait() + background_task = asyncio.create_task(background()) + await background_started.wait() + return "done" + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), sibling())) + + run_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + release_sibling.set() + paused = await asyncio.wait_for(run_task, timeout=5) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert background_task is not None + assert background_task.cancelled() + assert background_finished.is_set() + + async def test_gather_return_exceptions_does_not_expose_interruption(self): + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + @built_workflow + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(ask(), return_exceptions=True)) + + paused = await wf.run(1) + resumed = await wf.run(responses={"y": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert paused.get_outputs() == [] + assert resumed.get_outputs() == [[5]] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_step_cannot_cache_swallowed_interruption(self): + calls = 0 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int | str: + nonlocal calls + calls += 1 + try: + return await ctx.request_info("y?", int, request_id="y") + except asyncio.CancelledError: + return "fallback" + + @built_workflow + async def wf(_: int) -> int | str: + return await ask() + + paused = await wf.run(1) + resumed = await wf.run(responses={"y": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert paused.get_outputs() == [] + assert resumed.get_outputs() == [5] + assert calls == 2 + + async def test_step_cannot_cache_descendant_swallowed_interruption(self): + calls = 0 + + async def child(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int | BaseException: + nonlocal calls + calls += 1 + return (await asyncio.gather(child(ctx), return_exceptions=True))[0] + + @built_workflow + async def wf(_: int) -> int | BaseException: + return await ask() + + paused = await wf.run(1) + resumed = await wf.run(responses={"y": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert paused.get_outputs() == [] + assert resumed.get_outputs() == [5] + assert calls == 2 + + async def test_step_after_swallowed_interruption_does_not_self_await(self): + after_calls = 0 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "after") + async def after() -> str: + nonlocal after_calls + after_calls += 1 + return "done" + + @built_workflow + async def wf(_: int) -> list[object]: + values: list[object] = list(await asyncio.gather(ask(), return_exceptions=True)) + values.append(await after()) + return values + + paused = await asyncio.wait_for(wf.run(1), timeout=1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert paused.get_outputs() == [] + assert after_calls == 0 + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="CancelledError pause is Python >= 3.11") + async def test_state_after_swallowed_interruption_is_not_committed(self): + @built_workflow + async def wf(_: int, ctx: RunContext) -> str: + with suppress(asyncio.CancelledError): + await ctx.request_info("continue?", int, request_id="r") + ctx.set_state("phase", "after") + return ctx.get_state("phase") + + paused = await wf.run(1) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert wf._last_state == {} # pyright: ignore[reportPrivateUsage] + + resumed = await wf.run(responses={"r": 1}) + assert resumed.get_outputs() == ["after"] + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="TaskGroup is Python >= 3.11") + async def test_task_group_hitl_interrupt_is_handled(self): + storage = InMemoryCheckpointStorage() + sibling_started = asyncio.Event() + interrupting = asyncio.Event() + release_sibling = asyncio.Event() + sibling_cancelled = False + sibling_calls = 0 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await sibling_started.wait() + interrupting.set() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "sibling") + async def sibling() -> str: + nonlocal sibling_calls, sibling_cancelled + sibling_calls += 1 + sibling_started.set() + try: + await release_sibling.wait() + except asyncio.CancelledError: + sibling_cancelled = True + raise + return f"answer #{sibling_calls}" + + @built_workflow(checkpoint_storage=storage) + async def wf(_: int) -> list[object]: + task_group_type = cast(Any, asyncio).TaskGroup + async with task_group_type() as group: + ask_task = group.create_task(ask()) + sibling_task = group.create_task(sibling()) + return [ask_task.result(), sibling_task.result()] + + initial_task = asyncio.ensure_future(wf.run(1)) + await interrupting.wait() + await asyncio.sleep(0) + assert not initial_task.done() + assert not sibling_cancelled + release_sibling.set() + + paused = await initial_task + checkpoints = await storage.list_checkpoints(workflow_name="wf") + pause_checkpoint = checkpoints[-1] + resumed = await wf.run(responses={"y": 5}) + + assert paused.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + assert "answer #1" in pause_checkpoint.state["_step_cache"].values() + assert resumed.get_outputs() == [[5, "answer #1"]] + assert sibling_calls == 1 + + @pytest.mark.skipif(not hasattr(asyncio, "TaskGroup"), reason="TaskGroup is Python >= 3.11") + async def test_task_group_real_failure_excludes_internal_interrupt(self): + started = 0 + all_started = asyncio.Event() + release = asyncio.Event() + + async def rendezvous() -> None: + nonlocal started + started += 1 + if started == 2: + all_started.set() + await release.wait() + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await rendezvous() + return await ctx.request_info("y?", int, request_id="y") + + @step(replay_key=lambda: "fail") + async def fail() -> None: + await rendezvous() + raise ValueError("boom") + + @built_workflow + async def wf(_: int) -> None: + task_group_type = cast(Any, asyncio).TaskGroup + async with task_group_type() as group: + group.create_task(ask()) + group.create_task(fail()) + await all_started.wait() + release.set() + + with pytest.raises(Exception) as exc_info: + await wf.run(1) + + assert "boom" in repr(exc_info.value) + assert "WorkflowInterrupted" not in repr(exc_info.value) + def test_replay_state_helpers_restore_and_clear_the_full_bundle(self): """Replay helpers keep message, caches, state, and pending IDs together."""