Coverage for async_durable_execution/_operation/wait_for_condition.py: 90%
142 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +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 .._core import (
13 CallableRuntimeError,
14 Duration,
15 DurableContext,
16 ErrorObject,
17 ExecutionError,
18 ExecutionState,
19 InvocationError,
20 Operation,
21 OperationIdentifier,
22 OperationStatus,
23 OperationSubType,
24 OperationType,
25 OperationUpdate,
26 SerDes,
27 ValidationError,
28 _DelayStrategy,
29 _encode_sdk_control_error_data,
30 _register_sdk_control_error_type,
31 _restore_sdk_control_error,
32 bind_current_context,
33 create_eager_task,
34 duration_to_seconds,
35 get_current_context,
36 get_durable_context,
37 suspend_with_optional_resume_delay,
38 suspend_with_optional_resume_timestamp,
39)
40from .._primitive.base import OperationExecutor
41from .step import StepContext, get_step_context
42from ..extension import ExtensionStepResult, get_extension_context
44if TYPE_CHECKING:
45 from collections.abc import Awaitable
48T = TypeVar("T")
50logger = logging.getLogger(__name__)
53PollingStrategyFunction = Callable[[T, int], Duration | None]
54_LEGACY_WAIT_FOR_CONDITION_ERROR_TYPE_NAMES = (
55 "async_durable_execution._extension.wait_for_condition.WaitForConditionError",
56 "async_durable_execution.exceptions.WaitForConditionError",
57 "async_durable_execution.extension.wait_for_condition.WaitForConditionError",
58)
61class WaitForConditionError(ExecutionError):
62 """Raised when a wait_for_condition operation exhausts its attempts."""
65def _restore_wait_for_condition_error(
66 message: str,
67 _payload: str | None,
68) -> WaitForConditionError:
69 return WaitForConditionError(message)
72_register_sdk_control_error_type(
73 WaitForConditionError,
74 restore=_restore_wait_for_condition_error,
75 legacy_exception_type_names=_LEGACY_WAIT_FOR_CONDITION_ERROR_TYPE_NAMES,
76)
79@dataclass
80class PollingStrategy(_DelayStrategy, Generic[T]):
81 """Polling strategy for `wait_for_condition()`."""
83 def __call__(self, result: T, attempts_made: int) -> int | None:
84 """Return the next polling delay, or None to stop polling."""
85 if result:
86 return None
88 if attempts_made >= self.max_attempts:
89 msg = (
90 f"wait_for_condition exhausted {self.max_attempts} attempts "
91 "before the condition was met"
92 )
93 raise WaitForConditionError(msg)
95 return self.calculate_delay(attempts_made)
98class WaitForConditionOperationExecutor(OperationExecutor[T]):
99 """Compatibility executor for the pre-SPI private import path.
101 The public operation uses the shared stateful STEP executor. This class and
102 `_wait_for_condition` remain available through
103 `async_durable_execution._extension.wait_for_condition` so existing private
104 imports retain their original behavior during the package migration.
105 """
107 SERDES_OPERATION_TYPE = OperationType.STEP
109 def __init__(
110 self,
111 check: Callable[[T | None], Awaitable[T]],
112 initial_state: T | None,
113 state: ExecutionState,
114 operation_identifier: OperationIdentifier,
115 polling_strategy: PollingStrategyFunction[T] | None = None,
116 serdes: SerDes | None = None,
117 ) -> None:
118 """Initialize the wait_for_condition executor.
120 Args:
121 check: The check function to evaluate the condition
122 initial_state: The state to pass to the first condition evaluation
123 state: The execution state
124 operation_identifier: The operation identifier
125 polling_strategy: Optional strategy for deciding whether and when to poll
126 serdes: Optional serializer/deserializer for state payloads
127 """
128 super().__init__(state=state, operation_identifier=operation_identifier)
129 self.check = check
130 self.initial_state = initial_state
131 self.polling_strategy = polling_strategy
132 self.serdes = serdes
133 self.default_polling_strategy = PollingStrategy[T]()
135 async def start(self) -> T:
136 """Start a new wait_for_condition operation."""
137 start_operation = OperationUpdate.create_step_start(self.operation_identifier)
138 await self.create_checkpoint(start_operation, is_sync=False)
139 return await self.execute(None)
141 async def replay(self, operation: Operation) -> T:
142 """Replay an existing wait_for_condition operation from its checkpoint."""
143 if operation.status is OperationStatus.SUCCEEDED:
144 logger.debug(
145 "wait_for_condition already completed for id: %s, name: %s",
146 self.operation_identifier.operation_id,
147 self.operation_name,
148 )
149 result = (
150 operation.step_details.result
151 if operation.step_details is not None
152 else None
153 )
154 if result is None:
155 return cast("T", None)
156 return await self.deserialize_value(
157 data=result,
158 serdes=self.serdes,
159 operation=operation,
160 attempt=(
161 operation.step_details.attempt
162 if operation.step_details is not None
163 else None
164 ),
165 )
167 if operation.status is OperationStatus.FAILED:
168 error = (
169 operation.step_details.error
170 if operation.step_details is not None
171 else None
172 )
173 if error is None:
174 msg = (
175 "Unknown error. No ErrorObject exists on the Checkpoint Operation."
176 )
177 raise CallableRuntimeError(
178 message=msg,
179 error_type=None,
180 data=None,
181 stack_trace=None,
182 )
183 control_error = _restore_sdk_control_error(
184 error.message or "wait_for_condition failed",
185 error.type,
186 error.data,
187 )
188 if control_error is not None:
189 raise control_error
191 raise CallableRuntimeError.from_error_object(error)
193 if operation.status is OperationStatus.PENDING:
194 scheduled_timestamp = (
195 operation.step_details.next_attempt_timestamp
196 if operation.step_details is not None
197 else None
198 )
199 suspend_with_optional_resume_timestamp(
200 msg=f"wait_for_condition {self.operation_name or self.operation_identifier.operation_id} will retry at timestamp {scheduled_timestamp}",
201 datetime_timestamp=scheduled_timestamp,
202 )
204 if operation.status is not OperationStatus.STARTED:
205 start_operation = OperationUpdate.create_step_start(
206 self.operation_identifier
207 )
208 await self.create_checkpoint(start_operation, is_sync=False)
210 return await self.execute(operation)
212 async def execute(self, operation: Operation | None) -> T:
213 """Execute check function and handle decision.
215 Args:
216 operation: The checkpoint operation, if one exists.
218 Returns:
219 The final state when condition is met
221 Raises:
222 Suspends if condition not met
223 Raises error if check function fails
224 """
225 operation_details = operation.step_details if operation is not None else None
227 try:
228 # Determine current state from checkpoint
229 if (
230 operation is not None
231 and operation.status in {OperationStatus.STARTED, OperationStatus.READY}
232 and operation_details is not None
233 and operation_details.result is not None
234 ):
235 current_state = await self.deserialize_value(
236 data=operation_details.result,
237 serdes=self.serdes,
238 operation=operation,
239 attempt=operation_details.attempt,
240 )
241 else:
242 current_state = self.initial_state
244 # The checkpoint stores completed attempts, so the current attempt is one more.
245 attempt: int = 1
246 if operation_details is not None:
247 attempt = operation_details.attempt + 1
249 check_context = WaitForConditionCheckContext(
250 attempt=attempt,
251 execution_state=self.state,
252 operation_identifier=self.operation_identifier,
253 )
254 with bind_current_context(check_context):
255 new_state = await self.check(current_state)
257 serialized_state = await self.serialize_value(
258 value=new_state,
259 serdes=self.serdes,
260 attempt=attempt,
261 )
263 logger.debug(
264 "wait_for_condition check completed: %s, name: %s, attempt: %s",
265 self.operation_identifier.operation_id,
266 self.operation_name,
267 attempt,
268 )
270 suspend_delay_seconds = self._resolve_delay_seconds(new_state, attempt)
271 if suspend_delay_seconds is None:
272 success_operation = OperationUpdate.create_step_succeed(
273 self.operation_identifier,
274 serialized_state,
275 )
276 await self.create_checkpoint(success_operation)
278 logger.debug(
279 "✅ wait_for_condition stopped polling for id: %s, name: %s",
280 self.operation_identifier.operation_id,
281 self.operation_name,
282 )
283 else:
284 delay_seconds = suspend_delay_seconds
286 # We enforce a minimum delay second of 1, to match model behaviour.
287 if delay_seconds < 1:
288 logger.warning(
289 (
290 "wait_for_condition delay_seconds step for id: %s, name: %s,"
291 "is %d < 1. Setting to minimum of 1 seconds."
292 ),
293 self.operation_identifier.operation_id,
294 self.operation_identifier.name,
295 delay_seconds,
296 )
297 delay_seconds = 1
299 retry_operation = OperationUpdate.create_step_retry(
300 self.operation_identifier,
301 error=None,
302 payload=serialized_state,
303 next_attempt_delay_seconds=delay_seconds,
304 )
306 # Checkpoint RETRY operation with blocking (is_sync=True, default).
307 # Must ensure the current state and next attempt timestamp are persisted before suspending.
308 # This guarantees the polling state is durable and will resume correctly on the next invocation.
309 await self.create_checkpoint(retry_operation)
311 suspend_with_optional_resume_delay(
312 msg=f"wait_for_condition {self.operation_identifier.name or self.operation_identifier.operation_id} will retry in {suspend_delay_seconds} seconds",
313 delay_seconds=suspend_delay_seconds,
314 )
316 except Exception as e:
317 if isinstance(e, InvocationError) and e.is_retryable():
318 raise
320 # Mark as failed - waitForCondition doesn't have its own retry logic for errors
321 # If the check function throws, it's considered a failure
322 logger.exception(
323 "❌ wait_for_condition failed for id: %s, name: %s",
324 self.operation_identifier.operation_id,
325 self.operation_identifier.name,
326 )
328 error = ErrorObject.from_exception(e)
329 sdk_error_data = _encode_sdk_control_error_data(e)
330 if sdk_error_data is not None:
331 error = ErrorObject(
332 message=error.message,
333 type=error.type,
334 data=sdk_error_data,
335 stack_trace=error.stack_trace,
336 )
338 fail_operation = OperationUpdate.create_step_fail(
339 self.operation_identifier,
340 error,
341 )
342 # Checkpoint FAIL operation with blocking (is_sync=True, default).
343 # Must ensure the failure state is persisted before raising the exception.
344 # This guarantees the error is durable and the condition won't be re-evaluated on replay.
345 await self.create_checkpoint(fail_operation)
346 raise
348 # SUCCEED is already durable. Deserialization failures must propagate
349 # without attempting a terminal FAIL transition for the same operation.
350 return await self.deserialize_value(
351 data=serialized_state,
352 serdes=self.serdes,
353 attempt=attempt,
354 )
356 def _resolve_delay_seconds(self, new_state: T, attempt: int) -> int | None:
357 polling_strategy = self.polling_strategy or self.default_polling_strategy
358 wait_delay = polling_strategy(new_state, attempt)
360 if wait_delay is None:
361 return None
363 if isinstance(wait_delay, int | timedelta):
364 return duration_to_seconds(wait_delay, "polling_strategy delay")
366 msg = "wait_for_condition polling_strategy must return int seconds, timedelta, or None"
367 raise ValidationError(msg)
370def wait_for_condition(
371 check: Callable[[T | None], Awaitable[T]],
372 *,
373 initial_state: T | None = None,
374 name: str | None = None,
375 polling_strategy: PollingStrategyFunction[T] | None = None,
376 serdes: SerDes | None = None,
377) -> asyncio.Task[T]:
378 """Poll durable state until the configured strategy decides to stop waiting.
380 The check receives the current state, beginning with `initial_state`,
381 and returns the next state. The polling strategy receives that result and
382 returns the next polling delay, or None to stop polling and complete with
383 the latest result.
384 """
385 strategy = polling_strategy or PollingStrategy[T]()
387 async def check_attempt(state: T | None) -> ExtensionStepResult[T]:
388 step_context = get_step_context()
389 attempt = step_context.attempt or 1
390 check_context = WaitForConditionCheckContext(
391 attempt=attempt,
392 execution_state=step_context.execution_state,
393 operation_identifier=step_context.operation_identifier,
394 )
395 with bind_current_context(check_context):
396 new_state = await check(state)
398 wait_delay = strategy(new_state, attempt)
399 if wait_delay is None:
400 return ExtensionStepResult.succeed(new_state)
401 if not isinstance(wait_delay, int | timedelta):
402 msg = (
403 "wait_for_condition polling_strategy must return int seconds, "
404 "timedelta, or None"
405 )
406 raise ValidationError(msg)
408 delay_seconds = duration_to_seconds(
409 wait_delay,
410 "polling_strategy delay",
411 )
412 if delay_seconds < 1:
413 logger.warning(
414 (
415 "wait_for_condition delay_seconds for name %s is %d < 1. "
416 "Setting to minimum of 1 second."
417 ),
418 name,
419 delay_seconds,
420 )
421 return ExtensionStepResult.retry(new_state, wait_delay)
423 return (
424 get_extension_context()
425 ._reserve_sdk_operation(name) # noqa: SLF001
426 ._run_stateful_step( # noqa: SLF001
427 check_attempt,
428 sub_type=OperationSubType.WAIT_FOR_CONDITION,
429 initial_state=initial_state,
430 serdes=serdes,
431 raise_original_error=True,
432 )
433 )
436async def _wait_for_condition(
437 check: Callable[[T | None], Awaitable[T]],
438 *,
439 context: DurableContext,
440 operation_identifier: OperationIdentifier,
441 initial_state: T | None = None,
442 polling_strategy: PollingStrategyFunction[T] | None = None,
443 serdes: SerDes | None = None,
444) -> T:
445 """Run the compatibility executor retained for former private imports."""
446 executor: WaitForConditionOperationExecutor[T] = WaitForConditionOperationExecutor(
447 check=check,
448 initial_state=initial_state,
449 state=context.execution_state,
450 operation_identifier=operation_identifier,
451 polling_strategy=polling_strategy,
452 serdes=serdes,
453 )
454 return await executor.process()
457@dataclass(frozen=True)
458class WaitForConditionCheckContext(StepContext):
459 """Context available during wait_for_condition checker execution."""
461 pass
464def get_wait_for_condition_check_context() -> WaitForConditionCheckContext:
465 """Return the active `WaitForConditionCheckContext`."""
466 current_context = get_current_context()
467 if not isinstance(current_context, WaitForConditionCheckContext):
468 msg = (
469 "get_wait_for_condition_check_context() can only be used while a "
470 "wait_for_condition check is executing."
471 )
472 raise RuntimeError(msg)
473 return current_context