feat: add durable_wait_for_callback decorator

Add a new decorator that enables cleaner API for wait_for_callback
operations by allowing submitter functions to accept additional
parameters that are bound at call time.

The decorator wraps callables that take callback_id, context, and
additional parameters, returning a function that binds those extra
parameters and produces a submitter compatible with wait_for_callback.

Changes:
* Add durable_wait_for_callback decorator to context module
* Export decorator from package __init__ for public API
* Add integration test validating parameter binding and name propagation
* Verify operation names are correctly set on STEP and CALLBACK operations

Example usage:
  @durable_wait_for_callback
  def submit_task(callback_id, context, task_name, priority):
      external_api.submit(task_name, priority, callback_id)

  result = context.wait_for_callback(submit_task(my_task, priority=5))

Author: yaythomas

src/aws_durable_execution_sdk_python/__init__.py

@@ -5,6 +5,7 @@
 from aws_durable_execution_sdk_python.context import (
     DurableContext,
     durable_step,
+    durable_wait_for_callback,
     durable_with_child_context,
 )
 
@@ -30,5 +31,6 @@
     "ValidationError",
     "durable_execution",
     "durable_step",
+    "durable_wait_for_callback",
     "durable_with_child_context",
 ]

src/aws_durable_execution_sdk_python/context.py

@@ -104,6 +104,52 @@ def function_with_arguments(child_context: DurableContext):
     return wrapper
 
 
+def durable_wait_for_callback(
+    func: Callable[Concatenate[str, WaitForCallbackContext, Params], T],
+) -> Callable[Params, Callable[[str, WaitForCallbackContext], T]]:
+    """Wrap your callable into a wait_for_callback submitter function.
+
+    This decorator allows you to define a submitter function with additional
+    parameters that will be bound when called.
+
+    Args:
+        func: A callable that takes callback_id, context, and additional parameters
+
+    Returns:
+        A wrapper function that binds the additional parameters and returns
+        a submitter function compatible with wait_for_callback
+
+    Example:
+        @durable_wait_for_callback
+        def submit_to_external_system(
+            callback_id: str,
+            context: WaitForCallbackContext,
+            task_name: str,
+            priority: int
+        ):
+            context.logger.info(f"Submitting {task_name} with callback {callback_id}")
+            external_api.submit_task(
+                task_name=task_name,
+                priority=priority,
+                callback_id=callback_id
+            )
+
+        # Usage in durable handler:
+        result = context.wait_for_callback(
+            submit_to_external_system("my_task", priority=5)
+        )
+    """
+
+    def wrapper(*args, **kwargs):
+        def submitter_with_arguments(callback_id: str, context: WaitForCallbackContext):
+            return func(callback_id, context, *args, **kwargs)
+
+        submitter_with_arguments._original_name = func.__name__  # noqa: SLF001
+        return submitter_with_arguments
+
+    return wrapper
+
+
 class Callback(Generic[T], CallbackProtocol[T]):  # noqa: PYI059
     """A future that will block on result() until callback_id returns."""
 

tests/e2e/execution_int_test.py

@@ -11,18 +11,20 @@
 from aws_durable_execution_sdk_python.context import (
     DurableContext,
     durable_step,
+    durable_wait_for_callback,
     durable_with_child_context,
 )
 from aws_durable_execution_sdk_python.execution import (
     InvocationStatus,
     durable_execution,
 )
-
-# LambdaContext no longer needed - using duck typing
 from aws_durable_execution_sdk_python.lambda_service import (
+    CallbackDetails,
     CheckpointOutput,
     CheckpointUpdatedExecutionState,
+    Operation,
     OperationAction,
+    OperationStatus,
     OperationType,
 )
 from aws_durable_execution_sdk_python.logger import LoggerInterface
@@ -487,3 +489,120 @@ def mock_checkpoint(
         assert checkpoint.action is OperationAction.START
         assert checkpoint.operation_id == next(operation_ids)
         assert checkpoint.wait_options.wait_seconds == 1
+
+
+def test_durable_wait_for_callback_decorator():
+    """Test the durable_wait_for_callback decorator with additional parameters."""
+
+    mock_submitter = Mock()
+
+    @durable_wait_for_callback
+    def submit_to_external_system(callback_id, context, task_name, priority):
+        mock_submitter(callback_id, task_name, priority)
+        context.logger.info("Submitting %s with callback %s", task_name, callback_id)
+
+    @durable_execution
+    def my_handler(event, context):
+        context.wait_for_callback(submit_to_external_system("my_task", priority=5))
+
+    with patch(
+        "aws_durable_execution_sdk_python.execution.LambdaClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_client_class.initialize_from_env.return_value = mock_client
+
+        checkpoint_calls = []
+
+        def mock_checkpoint(
+            durable_execution_arn,
+            checkpoint_token,
+            updates,
+            client_token="token",  # noqa: S107
+        ):
+            checkpoint_calls.append(updates)
+
+            # For CALLBACK operations, return the operation with callback details
+            operations = [
+                Operation(
+                    operation_id=update.operation_id,
+                    operation_type=OperationType.CALLBACK,
+                    status=OperationStatus.STARTED,
+                    callback_details=CallbackDetails(
+                        callback_id=f"callback-{update.operation_id[:8]}"
+                    ),
+                )
+                for update in updates
+                if update.operation_type == OperationType.CALLBACK
+            ]
+
+            return CheckpointOutput(
+                checkpoint_token="new_token",  # noqa: S106
+                new_execution_state=CheckpointUpdatedExecutionState(
+                    operations=operations, next_marker=None
+                ),
+            )
+
+        mock_client.checkpoint = mock_checkpoint
+
+        event = {
+            "DurableExecutionArn": "test-arn",
+            "CheckpointToken": "test-token",
+            "InitialExecutionState": {
+                "Operations": [
+                    {
+                        "Id": "execution-1",
+                        "Type": "EXECUTION",
+                        "Status": "STARTED",
+                        "ExecutionDetails": {"InputPayload": "{}"},
+                    }
+                ],
+                "NextMarker": "",
+            },
+            "LocalRunner": True,
+        }
+
+        lambda_context = Mock()
+        lambda_context.aws_request_id = "test-request-id"
+        lambda_context.client_context = None
+        lambda_context.identity = None
+        lambda_context._epoch_deadline_time_in_ms = 0  # noqa: SLF001
+        lambda_context.invoked_function_arn = "test-arn"
+        lambda_context.tenant_id = None
+
+        result = my_handler(event, lambda_context)
+
+        assert result["Status"] == InvocationStatus.PENDING.value
+
+        all_operations = [op for batch in checkpoint_calls for op in batch]
+        assert len(all_operations) == 4
+
+        # First: CONTEXT START
+        first_checkpoint = all_operations[0]
+        assert first_checkpoint.operation_type is OperationType.CONTEXT
+        assert first_checkpoint.action is OperationAction.START
+        assert first_checkpoint.name == "submit_to_external_system"
+
+        # Second: CALLBACK START
+        second_checkpoint = all_operations[1]
+        assert second_checkpoint.operation_type is OperationType.CALLBACK
+        assert second_checkpoint.action is OperationAction.START
+        assert second_checkpoint.parent_id == first_checkpoint.operation_id
+        assert second_checkpoint.name == "submit_to_external_system create callback id"
+
+        # Third: STEP START
+        third_checkpoint = all_operations[2]
+        assert third_checkpoint.operation_type is OperationType.STEP
+        assert third_checkpoint.action is OperationAction.START
+        assert third_checkpoint.parent_id == first_checkpoint.operation_id
+        assert third_checkpoint.name == "submit_to_external_system submitter"
+
+        # Fourth: STEP SUCCEED
+        fourth_checkpoint = all_operations[3]
+        assert fourth_checkpoint.operation_type is OperationType.STEP
+        assert fourth_checkpoint.action is OperationAction.SUCCEED
+        assert fourth_checkpoint.operation_id == third_checkpoint.operation_id
+
+        mock_submitter.assert_called_once()
+        call_args = mock_submitter.call_args[0]
+        assert call_args[1] == "my_task"
+        assert call_args[2] == 5