Coverage for async-durable-execution/src/async_durable_execution/composite/wait_for_condition.py: 100%
113 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1"""Implement the durable wait_for_condition operation."""
3from __future__ import annotations
5import asyncio
6import logging
7from collections.abc import Callable
8from dataclasses import dataclass
9from datetime import timedelta
10from typing import TYPE_CHECKING, Generic, TypeVar, cast
12from ..config import _DelayStrategy, Duration, duration_to_seconds
13from ..context import (
14 bind_current_context,
15)
16from ..exceptions import (
17 CallableRuntimeError,
18 ExecutionError,
19 ValidationError,
20 suspend_with_optional_resume_delay,
21 suspend_with_optional_resume_timestamp,
22)
23from ..models import (
24 ErrorObject,
25 Operation,
26 OperationIdentifier,
27 OperationStatus,
28 OperationUpdate,
29 OperationSubType,
30)
31from ..primitive.base import OperationExecutor
32from ..primitive.child import get_durable_context
33from ..primitive.step import StepContext
34from ..task import create_eager_task
36if TYPE_CHECKING:
37 from collections.abc import Awaitable
39 from ..primitive.child import DurableContext
40 from ..serdes import SerDes
41 from ..state import ExecutionState
44T = TypeVar("T")
46logger = logging.getLogger(__name__)
49PollingStrategyFunction = Callable[[T, int], Duration | None]
52@dataclass
53class PollingStrategy(_DelayStrategy, Generic[T]):
54 """Polling strategy for `wait_for_condition()`."""
56 def __call__(self, result: T, attempts_made: int) -> int | None:
57 """Return the next polling delay, or None to stop polling."""
58 if result or attempts_made >= self.max_attempts:
59 return None
61 return self.calculate_delay(attempts_made)
64class WaitForConditionOperationExecutor(OperationExecutor[T]):
65 """Executor for wait_for_condition operations."""
67 def __init__(
68 self,
69 check: Callable[[T | None], Awaitable[T]],
70 initial_state: T | None,
71 state: ExecutionState,
72 operation_identifier: OperationIdentifier,
73 polling_strategy: PollingStrategyFunction[T] | None = None,
74 serdes: SerDes | None = None,
75 ):
76 """Initialize the wait_for_condition executor.
78 Args:
79 check: The check function to evaluate the condition
80 initial_state: The state to pass to the first condition evaluation
81 state: The execution state
82 operation_identifier: The operation identifier
83 polling_strategy: Optional strategy for deciding whether and when to poll
84 serdes: Optional serializer/deserializer for state payloads
85 """
86 super().__init__(state=state, operation_identifier=operation_identifier)
87 self.check = check
88 self.initial_state = initial_state
89 self.polling_strategy = polling_strategy
90 self.serdes = serdes
91 self.default_polling_strategy = PollingStrategy[T]()
93 async def start(self) -> T:
94 """Start a new wait_for_condition operation."""
95 start_operation = OperationUpdate.create_wait_for_condition_start(
96 identifier=self.operation_identifier,
97 )
98 await self.create_checkpoint(start_operation, is_sync=False)
99 return await self.execute(None)
101 async def replay(self, operation: Operation) -> T:
102 """Replay an existing wait_for_condition operation from its checkpoint."""
103 if operation.status is OperationStatus.SUCCEEDED:
104 logger.debug(
105 "wait_for_condition already completed for id: %s, name: %s",
106 self.operation_identifier.operation_id,
107 self.operation_name,
108 )
109 result = (
110 operation.step_details.result
111 if operation.step_details is not None
112 else None
113 )
114 if result is None:
115 return cast("T", None)
116 return await self.deserialize_value(
117 data=result,
118 serdes=self.serdes,
119 )
121 if operation.status is OperationStatus.FAILED:
122 error = (
123 operation.step_details.error
124 if operation.step_details is not None
125 else None
126 )
127 if error is None:
128 msg = (
129 "Unknown error. No ErrorObject exists on the Checkpoint Operation."
130 )
131 raise CallableRuntimeError(
132 message=msg,
133 error_type=None,
134 data=None,
135 stack_trace=None,
136 )
137 raise CallableRuntimeError.from_error_object(error)
139 if operation.status is OperationStatus.PENDING:
140 scheduled_timestamp = (
141 operation.step_details.next_attempt_timestamp
142 if operation.step_details is not None
143 else None
144 )
145 suspend_with_optional_resume_timestamp(
146 msg=f"wait_for_condition {self.operation_name or self.operation_identifier.operation_id} will retry at timestamp {scheduled_timestamp}",
147 datetime_timestamp=scheduled_timestamp,
148 )
150 if operation.status is not OperationStatus.STARTED:
151 start_operation = OperationUpdate.create_wait_for_condition_start(
152 identifier=self.operation_identifier,
153 )
154 await self.create_checkpoint(start_operation, is_sync=False)
156 return await self.execute(operation)
158 async def execute(self, operation: Operation | None) -> T:
159 """Execute check function and handle decision.
161 Args:
162 operation: The checkpoint operation, if one exists.
164 Returns:
165 The final state when condition is met
167 Raises:
168 Suspends if condition not met
169 Raises error if check function fails
170 """
171 # Determine current state from checkpoint
172 operation_details = operation.step_details if operation is not None else None
173 if (
174 operation is not None
175 and operation.status in {OperationStatus.STARTED, OperationStatus.READY}
176 and operation_details is not None
177 and operation_details.result
178 ):
179 try:
180 current_state = await self.deserialize_value(
181 data=operation_details.result,
182 serdes=self.serdes,
183 )
184 except Exception:
185 # Default to initial state if there's an error getting checkpointed state
186 logger.exception(
187 "⚠️ wait_for_condition failed to deserialize state for id: %s, name: %s. Using initial state.",
188 self.operation_identifier.operation_id,
189 self.operation_name,
190 )
191 current_state = self.initial_state
192 else:
193 current_state = self.initial_state
195 # Get attempt number - current attempt is checkpointed attempts + 1
196 # The checkpoint stores completed attempts, so the current attempt being executed is one more
197 attempt: int = 1
198 if operation_details is not None:
199 attempt = operation_details.attempt + 1
201 try:
202 check_context = WaitForConditionCheckContext(
203 attempt=attempt,
204 execution_state=self.state,
205 operation_identifier=self.operation_identifier,
206 )
207 with bind_current_context(check_context):
208 new_state = await self.check(current_state)
210 serialized_state = await self.serialize_value(
211 value=new_state,
212 serdes=self.serdes,
213 )
215 logger.debug(
216 "wait_for_condition check completed: %s, name: %s, attempt: %s",
217 self.operation_identifier.operation_id,
218 self.operation_name,
219 attempt,
220 )
222 suspend_delay_seconds = self._resolve_delay_seconds(new_state, attempt)
223 if suspend_delay_seconds is None:
224 success_operation = OperationUpdate.create_wait_for_condition_succeed(
225 identifier=self.operation_identifier,
226 payload=serialized_state,
227 )
228 await self.create_checkpoint(success_operation)
230 logger.debug(
231 "✅ wait_for_condition stopped polling for id: %s, name: %s",
232 self.operation_identifier.operation_id,
233 self.operation_name,
234 )
235 return await self.deserialize_value( # noqa: TRY300
236 data=serialized_state,
237 serdes=self.serdes,
238 )
240 delay_seconds = suspend_delay_seconds
242 # We enforce a minimum delay second of 1, to match model behaviour.
243 if delay_seconds < 1:
244 logger.warning(
245 (
246 "wait_for_condition delay_seconds step for id: %s, name: %s,"
247 "is %d < 1. Setting to minimum of 1 seconds."
248 ),
249 self.operation_identifier.operation_id,
250 self.operation_identifier.name,
251 delay_seconds,
252 )
253 delay_seconds = 1
255 retry_operation = OperationUpdate.create_wait_for_condition_retry(
256 identifier=self.operation_identifier,
257 payload=serialized_state,
258 next_attempt_delay_seconds=delay_seconds,
259 )
261 # Checkpoint RETRY operation with blocking (is_sync=True, default).
262 # Must ensure the current state and next attempt timestamp are persisted before suspending.
263 # This guarantees the polling state is durable and will resume correctly on the next invocation.
264 await self.create_checkpoint(retry_operation)
266 suspend_with_optional_resume_delay(
267 msg=f"wait_for_condition {self.operation_identifier.name or self.operation_identifier.operation_id} will retry in {suspend_delay_seconds} seconds",
268 delay_seconds=suspend_delay_seconds,
269 )
271 except Exception as e:
272 # Mark as failed - waitForCondition doesn't have its own retry logic for errors
273 # If the check function throws, it's considered a failure
274 logger.exception(
275 "❌ wait_for_condition failed for id: %s, name: %s",
276 self.operation_identifier.operation_id,
277 self.operation_identifier.name,
278 )
280 fail_operation = OperationUpdate.create_wait_for_condition_fail(
281 identifier=self.operation_identifier,
282 error=ErrorObject.from_exception(e),
283 )
284 # Checkpoint FAIL operation with blocking (is_sync=True, default).
285 # Must ensure the failure state is persisted before raising the exception.
286 # This guarantees the error is durable and the condition won't be re-evaluated on replay.
287 await self.create_checkpoint(fail_operation)
288 raise
290 msg: str = (
291 "wait_for_condition should never reach this point" # pragma: no cover
292 )
293 raise ExecutionError(msg) # pragma: no cover
295 def _resolve_delay_seconds(self, new_state: T, attempt: int) -> int | None:
296 polling_strategy = self.polling_strategy or self.default_polling_strategy
297 wait_delay = polling_strategy(new_state, attempt)
299 if wait_delay is None:
300 return None
302 if isinstance(wait_delay, int | timedelta):
303 return duration_to_seconds(wait_delay, "polling_strategy delay")
305 msg = "wait_for_condition polling_strategy must return int seconds, timedelta, or None"
306 raise ValidationError(msg)
309def wait_for_condition(
310 check: Callable[[T | None], Awaitable[T]],
311 *,
312 initial_state: T | None = None,
313 name: str | None = None,
314 polling_strategy: PollingStrategyFunction[T] | None = None,
315 serdes: SerDes | None = None,
316) -> asyncio.Task[T]:
317 """Poll durable state until the configured strategy decides to stop waiting.
319 The check receives the current state, beginning with `initial_state`,
320 and returns the next state. The polling strategy receives that result and
321 returns the next polling delay, or None to stop polling and complete with
322 the latest result.
323 """
324 context = get_durable_context("wait_for_condition")
326 with context._replay_aware(executes_user_code=True):
327 operation_id = context.step_counter.create_step_id()
328 operation_identifier = OperationIdentifier(
329 operation_id=operation_id,
330 sub_type=OperationSubType.WAIT_FOR_CONDITION,
331 parent_id=context.parent_id,
332 name=name,
333 )
335 return create_eager_task(
336 lambda: _wait_for_condition(
337 check=check,
338 context=context,
339 operation_identifier=operation_identifier,
340 initial_state=initial_state,
341 polling_strategy=polling_strategy,
342 serdes=serdes,
343 ),
344 )
347async def _wait_for_condition(
348 check: Callable[[T | None], Awaitable[T]],
349 *,
350 context: DurableContext,
351 operation_identifier: OperationIdentifier,
352 initial_state: T | None = None,
353 polling_strategy: PollingStrategyFunction[T] | None = None,
354 serdes: SerDes | None = None,
355) -> T:
356 executor: WaitForConditionOperationExecutor[T] = WaitForConditionOperationExecutor(
357 check=check,
358 initial_state=initial_state,
359 state=context.execution_state,
360 operation_identifier=operation_identifier,
361 polling_strategy=polling_strategy,
362 serdes=serdes,
363 )
364 return await executor.process()
367@dataclass(frozen=True)
368class WaitForConditionCheckContext(StepContext):
369 """Context available during wait_for_condition checker execution."""
371 pass