Coverage for async_durable_execution/extension.py: 84%
208 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1"""Stable contracts for authoring third-party durable operations."""
3from __future__ import annotations
5import asyncio
6from collections.abc import Awaitable, Callable, Coroutine
7from dataclasses import dataclass
8from typing import Any, Generic, TypeAlias, TypeVar, cast
10from ._core import (
11 Duration,
12 DurableContext,
13 LambdaContext,
14 OperationIdentifier,
15 OperationSubType,
16 OperationSubTypeValue,
17 OperationType,
18 SerDes,
19 ValidationError,
20 create_eager_task,
21 duration_to_seconds,
22 get_durable_context,
23)
24from ._primitive.callback import Callback, _create_callback
25from ._primitive.child import SummaryGenerator, _run_child_context
26from ._primitive.invoke import _invoke
27from ._primitive.step import (
28 StepSemantics,
29 _stateful_step,
30 _step,
31)
32from ._primitive.wait import _wait
34T = TypeVar("T")
35P = TypeVar("P")
36R = TypeVar("R")
37U = TypeVar("U")
40@dataclass(frozen=True)
41class ExtensionStepResult(Generic[T]):
42 """Outcome returned by one attempt of a stateful extension step."""
44 value: T
45 retry_delay: Duration | None = None
47 @classmethod
48 def succeed(cls, value: T) -> ExtensionStepResult[T]:
49 """Complete the extension step with ``value``."""
50 return cls(value=value)
52 @classmethod
53 def retry(cls, state: T | None, delay: Duration) -> ExtensionStepResult[T]:
54 """Checkpoint ``state`` and retry after ``delay``."""
55 duration_to_seconds(delay, "retry delay")
56 return cls(value=cast("T", state), retry_delay=delay)
58 @property
59 def is_retry(self) -> bool:
60 """Return whether this outcome schedules another attempt."""
61 return self.retry_delay is not None
64ExtensionStepFunction: TypeAlias = Callable[
65 [T | None],
66 Awaitable[ExtensionStepResult[T]],
67]
68ExtensionStepRetryStrategy: TypeAlias = Callable[
69 [Exception, T | None, int],
70 ExtensionStepResult[T] | None,
71]
74def _normalize_operation_name(name: str | None) -> str | None:
75 if name is None: 75 ↛ 76line 75 didn't jump to line 76 because the condition on line 75 was never true
76 return None
77 if not isinstance(name, str): 77 ↛ 78line 77 didn't jump to line 78 because the condition on line 77 was never true
78 msg = "name must be a string or None"
79 raise TypeError(msg)
80 if not name.strip():
81 msg = "name must not be blank"
82 raise ValueError(msg)
83 return name
86def _normalize_sdk_operation_name(name: str | None) -> str | None:
87 """Normalize names the service omits while preserving legacy SDK inputs."""
88 return None if name == "" else name
91def _normalize_sub_type(sub_type: str | OperationSubType) -> OperationSubTypeValue:
92 if isinstance(sub_type, OperationSubType):
93 return sub_type
94 if not isinstance(sub_type, str): 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true
95 msg = "sub_type must be a string or OperationSubType"
96 raise TypeError(msg)
97 if not sub_type.strip(): 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true
98 msg = "sub_type must not be blank"
99 raise ValueError(msg)
100 try:
101 reserved_sub_type = OperationSubType(sub_type)
102 except ValueError:
103 return sub_type
104 msg = (
105 f"sub_type {reserved_sub_type.value!r} is reserved by the SDK; "
106 "use an extension-owned subtype string"
107 )
108 raise ValueError(msg)
111class ExtensionOperation:
112 """Opaque one-shot reservation for one SDK-owned durable primitive.
114 Instances are created only by :meth:`ExtensionContext.reserve`.
115 """
117 __slots__ = (
118 "_claimed_operation_type",
119 "_claimed",
120 "_context",
121 "_identifier",
122 "_name",
123 "_operation_id",
124 "_parent_replaying",
125 "_replaying",
126 )
128 _claimed_operation_type: OperationType | None
129 _claimed: bool
130 _context: DurableContext
131 _identifier: OperationIdentifier | None
132 _name: str | None
133 _operation_id: str
134 _parent_replaying: bool
135 _replaying: bool
137 def __init__(self) -> None:
138 msg = "ExtensionOperation instances are created by ExtensionContext.reserve()"
139 raise TypeError(msg)
141 def step(
142 self,
143 func: ExtensionStepFunction[T],
144 *,
145 sub_type: str | OperationSubType,
146 initial_state: T | None = None,
147 retry_strategy: ExtensionStepRetryStrategy[T] | None = None,
148 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
149 serdes: SerDes[T] | None = None,
150 ) -> asyncio.Task[T]:
151 """Use this reservation for a stateful STEP primitive."""
152 identifier = self._claim(OperationType.STEP, sub_type)
153 return self._create_task(
154 lambda: _stateful_step(
155 func=func,
156 context=self._context,
157 operation_identifier=identifier,
158 initial_state=initial_state,
159 retry_strategy=retry_strategy,
160 step_semantics=step_semantics,
161 serdes=serdes,
162 ),
163 executes_user_code=True,
164 )
166 def _run_stateful_step(
167 self,
168 func: ExtensionStepFunction[T],
169 *,
170 sub_type: str | OperationSubType,
171 initial_state: T | None = None,
172 retry_strategy: ExtensionStepRetryStrategy[T] | None = None,
173 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
174 serdes: SerDes[T] | None = None,
175 raise_original_error: bool = False,
176 ) -> asyncio.Task[T]:
177 """Use this reservation for an SDK-owned stateful STEP primitive."""
178 identifier = self._claim(
179 OperationType.STEP,
180 sub_type,
181 include_operation_type=False,
182 )
183 return self._create_task(
184 lambda: _stateful_step(
185 func=func,
186 context=self._context,
187 operation_identifier=identifier,
188 initial_state=initial_state,
189 retry_strategy=retry_strategy,
190 step_semantics=step_semantics,
191 serdes=serdes,
192 raise_original_error=raise_original_error,
193 ),
194 executes_user_code=True,
195 )
197 def _run_step(
198 self,
199 func: Callable[[], Awaitable[T]],
200 *,
201 sub_type: str | OperationSubType,
202 retry_strategy: Callable[[Exception, int], Duration | None] | None = None,
203 step_semantics: StepSemantics = StepSemantics.AT_LEAST_ONCE_PER_RETRY,
204 serdes: SerDes[T] | None = None,
205 ) -> asyncio.Task[T]:
206 """Use this reservation for an SDK-owned standard STEP primitive."""
207 identifier = self._claim(
208 OperationType.STEP,
209 sub_type,
210 include_operation_type=False,
211 )
212 return self._create_task(
213 lambda: _step(
214 func=func,
215 context=self._context,
216 operation_identifier=identifier,
217 retry_strategy=retry_strategy,
218 step_semantics=step_semantics,
219 serdes=serdes,
220 ),
221 executes_user_code=True,
222 )
224 def wait(
225 self,
226 duration: Duration,
227 *,
228 sub_type: str | OperationSubType,
229 ) -> asyncio.Task[None]:
230 """Use this reservation for a WAIT primitive."""
231 seconds = duration_to_seconds(duration)
232 if seconds < 1: 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 msg = "duration must be at least 1 second"
234 raise ValidationError(msg)
235 identifier = self._claim(OperationType.WAIT, sub_type)
236 return self._create_task(
237 lambda: _wait(
238 seconds=seconds,
239 context=self._context,
240 operation_identifier=identifier,
241 ),
242 executes_user_code=False,
243 )
245 def _run_wait(
246 self,
247 duration: Duration,
248 *,
249 sub_type: str | OperationSubType,
250 ) -> asyncio.Task[None]:
251 """Use this reservation for an SDK-owned WAIT primitive."""
252 seconds = duration_to_seconds(duration)
253 if seconds < 1: 253 ↛ 254line 253 didn't jump to line 254 because the condition on line 253 was never true
254 msg = "duration must be at least 1 second"
255 raise ValidationError(msg)
256 identifier = self._claim(
257 OperationType.WAIT,
258 sub_type,
259 include_operation_type=False,
260 )
261 return self._create_task(
262 lambda: _wait(
263 seconds=seconds,
264 context=self._context,
265 operation_identifier=identifier,
266 ),
267 executes_user_code=False,
268 )
270 def invoke(
271 self,
272 function_name: str,
273 payload: P,
274 *,
275 sub_type: str | OperationSubType,
276 serdes_payload: SerDes[P] | None = None,
277 serdes_result: SerDes[R] | None = None,
278 tenant_id: str | None = None,
279 ) -> asyncio.Task[R]:
280 """Use this reservation for a CHAINED_INVOKE primitive."""
281 identifier = self._claim(OperationType.CHAINED_INVOKE, sub_type)
282 return self._create_task(
283 lambda: _invoke(
284 function_name=function_name,
285 payload=payload,
286 context=self._context,
287 operation_identifier=identifier,
288 serdes_payload=serdes_payload,
289 serdes_result=serdes_result,
290 tenant_id=tenant_id,
291 ),
292 executes_user_code=False,
293 )
295 def _run_invoke(
296 self,
297 function_name: str,
298 payload: P,
299 *,
300 sub_type: str | OperationSubType,
301 serdes_payload: SerDes[P] | None = None,
302 serdes_result: SerDes[R] | None = None,
303 tenant_id: str | None = None,
304 ) -> asyncio.Task[R]:
305 """Use this reservation for an SDK-owned CHAINED_INVOKE primitive."""
306 identifier = self._claim(
307 OperationType.CHAINED_INVOKE,
308 sub_type,
309 include_operation_type=False,
310 )
311 return self._create_task(
312 lambda: _invoke(
313 function_name=function_name,
314 payload=payload,
315 context=self._context,
316 operation_identifier=identifier,
317 serdes_payload=serdes_payload,
318 serdes_result=serdes_result,
319 tenant_id=tenant_id,
320 ),
321 executes_user_code=False,
322 )
324 def create_callback(
325 self,
326 *,
327 sub_type: str | OperationSubType,
328 timeout: Duration | None = None,
329 heartbeat_timeout: Duration | None = None,
330 serdes: SerDes[T] | None = None,
331 ) -> asyncio.Task[Callback[T]]:
332 """Use this reservation for a CALLBACK primitive."""
333 identifier = self._claim(OperationType.CALLBACK, sub_type)
334 return self._create_task(
335 lambda: _create_callback(
336 context=self._context,
337 operation_identifier=identifier,
338 operation_id=self._operation_id,
339 timeout=timeout,
340 heartbeat_timeout=heartbeat_timeout,
341 serdes=serdes,
342 ),
343 executes_user_code=False,
344 )
346 def _run_create_callback(
347 self,
348 *,
349 sub_type: str | OperationSubType,
350 timeout: Duration | None = None,
351 heartbeat_timeout: Duration | None = None,
352 serdes: SerDes[T] | None = None,
353 ) -> asyncio.Task[Callback[T]]:
354 """Use this reservation for an SDK-owned CALLBACK primitive."""
355 identifier = self._claim(
356 OperationType.CALLBACK,
357 sub_type,
358 include_operation_type=False,
359 )
360 return self._create_task(
361 lambda: _create_callback(
362 context=self._context,
363 operation_identifier=identifier,
364 operation_id=self._operation_id,
365 timeout=timeout,
366 heartbeat_timeout=heartbeat_timeout,
367 serdes=serdes,
368 ),
369 executes_user_code=False,
370 )
372 def run_in_child_context(
373 self,
374 func: Callable[[], Awaitable[T]],
375 *,
376 sub_type: str | OperationSubType,
377 serdes: SerDes[T] | None = None,
378 summary_generator: SummaryGenerator[T] | None = None,
379 is_virtual: bool = False,
380 ) -> asyncio.Task[T]:
381 """Use this reservation for a CONTEXT primitive."""
382 identifier = self._claim(OperationType.CONTEXT, sub_type)
383 return self._create_child_context_task(
384 identifier,
385 func,
386 serdes=serdes,
387 summary_generator=summary_generator,
388 is_virtual=is_virtual,
389 replay_aware=True,
390 )
392 def _run_in_child_context(
393 self,
394 func: Callable[[], Awaitable[T]],
395 *,
396 sub_type: str | OperationSubType,
397 serdes: SerDes[T] | None = None,
398 summary_generator: SummaryGenerator[T] | None = None,
399 is_virtual: bool = False,
400 ) -> asyncio.Task[T]:
401 """Use this reservation for an SDK-owned CONTEXT primitive."""
402 identifier = self._claim(
403 OperationType.CONTEXT,
404 sub_type,
405 include_operation_type=False,
406 )
407 return self._create_child_context_task(
408 identifier,
409 func,
410 serdes=serdes,
411 summary_generator=summary_generator,
412 is_virtual=is_virtual,
413 replay_aware=True,
414 )
416 def _restart_child_context(
417 self,
418 func: Callable[[], Awaitable[T]],
419 *,
420 serdes: SerDes[T] | None = None,
421 summary_generator: SummaryGenerator[T] | None = None,
422 is_virtual: bool = False,
423 ) -> asyncio.Task[T]:
424 """Re-enter an SDK-owned child operation after an in-process suspension."""
425 identifier = self._identifier
426 if (
427 identifier is None
428 or self._claimed_operation_type is not OperationType.CONTEXT
429 ):
430 msg = "Only a claimed child-context reservation can be restarted"
431 raise RuntimeError(msg)
432 return self._create_child_context_task(
433 identifier,
434 func,
435 serdes=serdes,
436 summary_generator=summary_generator,
437 is_virtual=is_virtual,
438 replaying=True,
439 )
441 def _create_child_context_task(
442 self,
443 identifier: OperationIdentifier,
444 func: Callable[[], Awaitable[T]],
445 *,
446 serdes: SerDes[T] | None,
447 summary_generator: SummaryGenerator[T] | None,
448 is_virtual: bool,
449 replay_aware: bool = False,
450 replaying: bool | None = None,
451 ) -> asyncio.Task[T]:
452 child_context = self._context.create_child_context(
453 operation_id=self._operation_id,
454 is_virtual=is_virtual,
455 replaying=(self._parent_replaying if is_virtual else self._replaying)
456 if replaying is None
457 else replaying,
458 )
460 async def execute_child_context() -> T:
461 return await _run_child_context(
462 func,
463 context=self._context,
464 child_context=child_context,
465 operation_identifier=identifier,
466 serdes=serdes,
467 summary_generator=summary_generator,
468 is_virtual=is_virtual,
469 )
471 async def run_child_context() -> T:
472 if not replay_aware: 472 ↛ 473line 472 didn't jump to line 473 because the condition on line 472 was never true
473 return await execute_child_context()
474 if is_virtual:
475 # Virtual contexts have no container checkpoint. Their nested
476 # operations refine the inherited replay snapshot independently.
477 try:
478 return await execute_child_context()
479 finally:
480 if (
481 self._context.is_replaying()
482 and not self._context._next_reserved_or_sequential_operation_exists() # noqa: SLF001
483 ):
484 self._context._set_replay_status_new() # noqa: SLF001
485 with self._context._replay_aware(
486 operation_id=self._operation_id,
487 consume_reservation=False,
488 ):
489 return await execute_child_context()
491 if replay_aware: 491 ↛ 495line 491 didn't jump to line 495 because the condition on line 491 was always true
492 self._context.step_counter._consume_reservation( # noqa: SLF001
493 self._operation_id
494 )
495 return create_eager_task(run_child_context)
497 def _create_task(
498 self,
499 coro_factory: Callable[[], Coroutine[Any, Any, U]],
500 *,
501 executes_user_code: bool,
502 ) -> asyncio.Task[U]:
503 with self._context._replay_aware(
504 operation_id=self._operation_id,
505 executes_user_code=executes_user_code,
506 ):
507 return create_eager_task(coro_factory)
509 def _claim(
510 self,
511 operation_type: OperationType,
512 sub_type: str | OperationSubType,
513 *,
514 include_operation_type: bool = True,
515 ) -> OperationIdentifier:
516 self._require_active_context()
517 normalized_sub_type = _normalize_sub_type(sub_type)
518 if self._claimed:
519 msg = "An extension operation reservation can only be used once"
520 raise RuntimeError(msg)
521 self._context.step_counter._mark_reservation_selected() # noqa: SLF001
522 self._claimed = True
523 self._claimed_operation_type = operation_type
524 self._identifier = OperationIdentifier(
525 operation_id=self._operation_id,
526 sub_type=normalized_sub_type,
527 parent_id=self._context.parent_id,
528 name=self._name,
529 operation_type=operation_type if include_operation_type else None,
530 )
531 return self._identifier
533 def _require_active_context(self) -> None:
534 current_context = get_durable_context()
535 if current_context is not self._context: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true
536 msg = (
537 "An extension operation reservation can only be used in the "
538 "durable context where it was created"
539 )
540 raise RuntimeError(msg)
543def _create_extension_operation(
544 context: DurableContext,
545 operation_id: str,
546 name: str | None,
547 *,
548 has_checkpoint: bool,
549 parent_replaying: bool | None = None,
550) -> ExtensionOperation:
551 operation = object.__new__(ExtensionOperation)
552 operation._context = context # noqa: SLF001
553 operation._operation_id = operation_id # noqa: SLF001
554 operation._name = name # noqa: SLF001
555 operation._parent_replaying = ( # noqa: SLF001
556 context._virtual_child_replay_snapshot() # noqa: SLF001
557 if parent_replaying is None
558 else parent_replaying
559 )
560 operation._replaying = has_checkpoint # noqa: SLF001
561 operation._claimed = False # noqa: SLF001
562 operation._claimed_operation_type = None # noqa: SLF001
563 operation._identifier = None # noqa: SLF001
564 return operation
567class ExtensionContext:
568 """Stable extension-author view of the current durable execution scope."""
570 __slots__ = ("_context",)
572 def __init__(self, context: DurableContext) -> None:
573 self._context = context
575 @classmethod
576 def get_current(cls) -> ExtensionContext:
577 """Return the extension context for the active handler or child scope."""
578 return cls(get_durable_context())
580 @property
581 def lambda_context(self) -> LambdaContext | None:
582 """Return the active AWS Lambda context, when available."""
583 return self._context.lambda_context
585 @property
586 def recursive_level(self) -> int:
587 """Return the durable self-invocation recursion level."""
588 return self._context.recursive_level
590 def is_replaying(self) -> bool:
591 """Return whether the current scope is replaying checkpointed work."""
592 return self._context.is_replaying()
594 def reserve(
595 self,
596 name: str | None = None,
597 *,
598 local_operation_id: str | None = None,
599 ) -> ExtensionOperation:
600 """Reserve a stable one-shot primitive identity.
602 Sequential reservations depend on deterministic reservation order.
603 A caller-provided local id remains stable when reservation order changes,
604 but it must be unique within the current durable context.
605 """
606 self._require_active_context()
607 return self._create_reservation(
608 _normalize_operation_name(name),
609 local_operation_id=local_operation_id,
610 )
612 def _reserve_sdk_operation(
613 self,
614 name: str | None = None,
615 *,
616 local_operation_id: str | None = None,
617 ) -> ExtensionOperation:
618 """Reserve an SDK-owned primitive without selecting its behavior."""
619 self._require_active_context()
620 return self._create_reservation(
621 _normalize_sdk_operation_name(name),
622 local_operation_id=local_operation_id,
623 )
625 def _reserve_sdk_operation_id(
626 self,
627 name: str | None,
628 *,
629 operation_id: str,
630 parent_replaying: bool | None = None,
631 ) -> ExtensionOperation:
632 """Reserve an SDK-owned primitive with a replay-compatible operation id."""
633 self._require_active_context()
634 return self._create_reservation_for_operation_id(
635 _normalize_sdk_operation_name(name),
636 operation_id=operation_id,
637 parent_replaying=parent_replaying,
638 )
640 def _create_reservation(
641 self,
642 name: str | None,
643 *,
644 local_operation_id: str | None,
645 ) -> ExtensionOperation:
646 operation_id = self._reserve_operation_id(local_operation_id)
647 return self._create_reservation_for_operation_id(
648 name,
649 operation_id=operation_id,
650 )
652 def _create_reservation_for_operation_id(
653 self,
654 name: str | None,
655 *,
656 operation_id: str,
657 parent_replaying: bool | None = None,
658 ) -> ExtensionOperation:
659 has_checkpoint = self._context._operation_result(operation_id) is not None # noqa: SLF001
660 self._context.step_counter._register_reservation( # noqa: SLF001
661 operation_id,
662 has_checkpoint=has_checkpoint,
663 )
664 return _create_extension_operation(
665 self._context,
666 operation_id,
667 name,
668 has_checkpoint=has_checkpoint,
669 parent_replaying=parent_replaying,
670 )
672 def _reserve_operation_id(self, local_operation_id: str | None) -> str:
673 if local_operation_id is None:
674 return self._context.step_counter.create_step_id()
675 return self._context.step_counter.create_step_id_for_local_id(
676 local_operation_id
677 )
679 def _require_active_context(self) -> None:
680 current_context = get_durable_context()
681 if current_context is not self._context: 681 ↛ 682line 681 didn't jump to line 682 because the condition on line 681 was never true
682 msg = (
683 "An extension context can only reserve operations in the "
684 "durable context where it was created"
685 )
686 raise RuntimeError(msg)
689def get_extension_context() -> ExtensionContext:
690 """Return the stable extension-author context for the active durable scope."""
691 return ExtensionContext.get_current()
694__all__ = [
695 "ExtensionContext",
696 "ExtensionOperation",
697 "ExtensionStepFunction",
698 "ExtensionStepResult",
699 "ExtensionStepRetryStrategy",
700 "get_extension_context",
701]