Coverage for async-durable-execution/src/async_durable_execution/primitive/callback.py: 100%
79 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1"""Implementation for the backend-supported create_callback operation."""
3from __future__ import annotations
5import asyncio
6import logging
7from typing import TYPE_CHECKING, Any, Generic, TypeVar
9from .base import OperationExecutor
10from .child import get_durable_context
11from ..config import Duration, duration_to_seconds
12from ..exceptions import ExecutionError, SuspendExecution, TerminationReason
13from ..models import (
14 CallbackOptions,
15 CallbackTimeoutType,
16 Operation,
17 OperationIdentifier,
18 OperationStatus,
19 OperationUpdate,
20 OperationSubType,
21)
22from ..serdes import deserialize, PassThroughSerDes
23from ..task import create_eager_task
25if TYPE_CHECKING:
26 from .child import DurableContext
27 from ..serdes import SerDes
28 from ..state import ExecutionState
30T = TypeVar("T") # Result type
32logger = logging.getLogger(__name__)
34PASS_THROUGH_SERDES: SerDes[Any] = PassThroughSerDes()
37class CallbackError(ExecutionError):
38 """Error in callback handling."""
40 def __init__(self, message: str, callback_id: str | None = None):
41 super().__init__(message, TerminationReason.CALLBACK_ERROR)
42 self.callback_id = callback_id
45class CallbackOperationExecutor(OperationExecutor[str]):
46 """Executor for callback operations."""
48 def __init__(
49 self,
50 state: ExecutionState,
51 operation_identifier: OperationIdentifier,
52 timeout: Duration | None = None,
53 heartbeat_timeout: Duration | None = None,
54 ):
55 """Initialize the callback operation executor.
57 Args:
58 state: The execution state
59 operation_identifier: The operation identifier
60 timeout: Optional maximum time to wait for callback completion.
61 heartbeat_timeout: Optional maximum time to wait between callback heartbeats.
62 """
63 super().__init__(state=state, operation_identifier=operation_identifier)
64 self.timeout_seconds = (
65 duration_to_seconds(timeout, "timeout") if timeout is not None else 0
66 )
67 self.heartbeat_timeout_seconds = (
68 duration_to_seconds(heartbeat_timeout, "heartbeat_timeout")
69 if heartbeat_timeout is not None
70 else 0
71 )
73 async def start(self) -> str:
74 """Start a new callback operation."""
75 callback_options = CallbackOptions(
76 timeout_seconds=self.timeout_seconds,
77 heartbeat_timeout_seconds=self.heartbeat_timeout_seconds,
78 )
80 create_callback_operation: OperationUpdate = OperationUpdate.create_callback(
81 identifier=self.operation_identifier,
82 callback_options=callback_options,
83 )
85 operation = await self.create_checkpoint(create_callback_operation)
87 if not operation:
88 msg = f"Missing callback details for operation: {self.operation_identifier.operation_id}"
89 raise CallbackError(msg)
90 return await self.replay(operation)
92 async def replay(self, operation: Operation) -> str:
93 """Replay an existing callback operation from its checkpoint."""
94 if not operation.callback_details:
95 msg = (
96 f"Missing callback details for operation: "
97 f"{self.operation_identifier.operation_id}"
98 )
99 raise CallbackError(msg)
101 return await self.execute(operation)
103 async def execute(self, operation: Operation | None) -> str:
104 """Execute callback operation by extracting the callback_id.
106 Callbacks don't execute logic - they just extract and return the callback_id
107 from the operation data.
109 Args:
110 operation: The callback operation containing callback_details
112 Returns:
113 The callback_id from the checkpoint
115 Raises:
116 CallbackError: If callback_details are missing (should never happen)
117 """
118 if operation is None or not operation.callback_details:
119 msg = f"Missing callback details for operation: {self.operation_identifier.operation_id}"
120 raise CallbackError(msg)
122 return operation.callback_details.callback_id
125def create_callback(
126 *,
127 name: str | None = None,
128 timeout: Duration | None = None,
129 heartbeat_timeout: Duration | None = None,
130 serdes: SerDes | None = None,
131) -> asyncio.Task[Callback]:
132 """Create a durable callback handle that external systems can complete later.
134 Args:
135 name: Optional durable operation name.
136 timeout: Optional maximum time to wait for callback completion.
137 heartbeat_timeout: Optional maximum time to wait between callback heartbeats.
138 serdes: Optional serializer for callback results.
139 """
140 context = get_durable_context("create_callback")
142 with context._replay_aware():
143 operation_id: str = context.step_counter.create_step_id()
144 operation_identifier = OperationIdentifier(
145 operation_id=operation_id,
146 sub_type=OperationSubType.CALLBACK,
147 parent_id=context.parent_id,
148 name=name,
149 )
151 return create_eager_task(
152 lambda: _create_callback(
153 context=context,
154 operation_identifier=operation_identifier,
155 operation_id=operation_id,
156 timeout=timeout,
157 heartbeat_timeout=heartbeat_timeout,
158 serdes=serdes,
159 ),
160 )
163async def _create_callback(
164 *,
165 context: DurableContext,
166 operation_identifier: OperationIdentifier,
167 operation_id: str,
168 timeout: Duration | None = None,
169 heartbeat_timeout: Duration | None = None,
170 serdes: SerDes | None = None,
171) -> Callback:
172 executor: CallbackOperationExecutor = CallbackOperationExecutor(
173 state=context.execution_state,
174 operation_identifier=operation_identifier,
175 timeout=timeout,
176 heartbeat_timeout=heartbeat_timeout,
177 )
178 callback_id: str = await executor.process()
179 return Callback(
180 callback_id=callback_id,
181 operation_id=operation_id,
182 state=context.execution_state,
183 serdes=serdes,
184 )
187class Callback(Generic[T]): # noqa: PYI059
188 """A future that will block on result() until callback_id returns."""
190 def __init__(
191 self,
192 callback_id: str,
193 operation_id: str,
194 state: ExecutionState,
195 serdes: SerDes[T] | None = None,
196 ):
197 self.callback_id: str = callback_id
198 self.operation_id: str = operation_id
199 self.state: ExecutionState = state
200 self.serdes: SerDes[T] | None = serdes
202 async def result(self) -> T | None:
203 """Return the result of the future. Will block until result is available.
205 This will suspend the current execution while waiting for the result to
206 become available. Durable Functions will replay the execution once the
207 result is ready, and proceed when it reaches the .result() call.
209 Use the callback id with the following APIs to send back the result, error or
210 heartbeats: SendDurableExecutionCallbackSuccess, SendDurableExecutionCallbackFailure
211 and SendDurableExecutionCallbackHeartbeat.
212 """
213 operation = self.state.operations.get(self.operation_id)
215 if not isinstance(operation, Operation):
216 msg = "Callback operation must exist"
217 raise CallbackError(message=msg, callback_id=self.callback_id)
219 if operation.status in {
220 OperationStatus.FAILED,
221 OperationStatus.CANCELLED,
222 OperationStatus.TIMED_OUT,
223 OperationStatus.STOPPED,
224 }:
225 msg = _format_callback_error_message(operation)
226 raise CallbackError(message=msg, callback_id=self.callback_id)
228 if operation.status is OperationStatus.SUCCEEDED:
229 if (
230 not operation.callback_details
231 or operation.callback_details.result is None
232 ):
233 return None
235 return await deserialize(
236 serdes=self.serdes if self.serdes is not None else PASS_THROUGH_SERDES,
237 data=operation.callback_details.result,
238 operation_id=self.operation_id,
239 durable_execution_arn=self.state.durable_execution_arn,
240 )
242 # operation exists; it has not terminated (successfully or otherwise)
243 # therefore we should wait
244 msg = "Callback result not received yet. Suspending execution while waiting for result."
245 raise SuspendExecution(msg)
248def _format_callback_error_message(operation: Operation) -> str:
249 """Build a stable callback error message from checkpoint state."""
250 error = operation.callback_details.error if operation.callback_details else None
251 if not error or not error.message:
252 return "Callback failed"
254 message = error.message
255 if (
256 operation.status is OperationStatus.TIMED_OUT
257 and error.type in {timeout.value for timeout in CallbackTimeoutType}
258 and error.type not in message
259 ):
260 return f"{message}: {error.type}"
262 return message