pass retry config to the Step that wraps the submitter

Author: Rares Polenciuc

src/aws_durable_execution_sdk_python/operation/callback.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING, Any
 
+from aws_durable_execution_sdk_python.config import StepConfig
 from aws_durable_execution_sdk_python.exceptions import FatalError
 from aws_durable_execution_sdk_python.lambda_service import (
     CallbackOptions,
@@ -16,7 +17,7 @@
     from aws_durable_execution_sdk_python.config import (
         CallbackConfig,
         WaitForCallbackConfig,
-    )
+)
     from aws_durable_execution_sdk_python.identifier import OperationIdentifier
     from aws_durable_execution_sdk_python.state import (
         CheckpointedResult,
@@ -97,6 +98,13 @@ def wait_for_callback_handler(
     def submitter_step(step_context):  # noqa: ARG001
         return submitter(callback.callback_id)
 
-    context.step(func=submitter_step, name=f"{name_with_space}submitter")
+    if config:
+        step_config = StepConfig(
+            retry_strategy=config.retry_strategy,
+            serdes=config.serdes,
+        )
+        context.step(func=submitter_step, name=f"{name_with_space}submitter", config=step_config)
+    else:
+        context.step(func=submitter_step, name=f"{name_with_space}submitter")
 
     return callback.result()

tests/operation/callback_test.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from aws_durable_execution_sdk_python.config import CallbackConfig
+from aws_durable_execution_sdk_python.config import CallbackConfig, WaitForCallbackConfig
 from aws_durable_execution_sdk_python.exceptions import FatalError
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -269,7 +269,7 @@ def test_wait_for_callback_handler_with_name_and_config():
     mock_callback.result.return_value = "named_callback_result"
     mock_context.create_callback.return_value = mock_callback
     mock_submitter = Mock()
-    config = CallbackConfig()
+    config = WaitForCallbackConfig()
 
     result = wait_for_callback_handler(
         mock_context, mock_submitter, "test_callback", config
@@ -291,7 +291,7 @@ def test_wait_for_callback_handler_submitter_called_with_callback_id():
     mock_context.create_callback.return_value = mock_callback
     mock_submitter = Mock()
 
-    def capture_step_call(func, name):
+    def capture_step_call(func, name, config=None):
         # Execute the step callable to verify submitter is called correctly
         step_context = Mock(spec=StepContext)
         func(step_context)
@@ -357,7 +357,7 @@ def test_wait_for_callback_handler_with_none_callback_id():
     mock_context.create_callback.return_value = mock_callback
     mock_submitter = Mock()
 
-    def execute_step(func, name):
+    def execute_step(func, name, config=None):
         step_context = Mock(spec=StepContext)
         return func(step_context)
 
@@ -378,7 +378,7 @@ def test_wait_for_callback_handler_with_empty_string_callback_id():
     mock_context.create_callback.return_value = mock_callback
     mock_submitter = Mock()
 
-    def execute_step(func, name):
+    def execute_step(func, name, config=None):
         step_context = Mock(spec=StepContext)
         return func(step_context)
 
@@ -591,7 +591,7 @@ def failing_submitter(callback_id):
         msg = "Submitter failed"
         raise ValueError(msg)
 
-    def step_side_effect(func, name):
+    def step_side_effect(func, name, config=None):
         step_context = Mock(spec=StepContext)
         func(step_context)
 
@@ -675,7 +675,7 @@ def test_wait_for_callback_handler_config_propagation():
     mock_context.create_callback.return_value = mock_callback
     mock_submitter = Mock()
 
-    config = CallbackConfig(timeout_seconds=120, heartbeat_timeout_seconds=30)
+    config = WaitForCallbackConfig(timeout_seconds=120, heartbeat_timeout_seconds=30)
 
     result = wait_for_callback_handler(
         mock_context, mock_submitter, "config_test", config
@@ -729,7 +729,7 @@ def test_callback_lifecycle_complete_flow():
     mock_callback.result.return_value = {"status": "completed", "data": "test_data"}
     mock_context.create_callback.return_value = mock_callback
 
-    config = CallbackConfig(timeout_seconds=300, heartbeat_timeout_seconds=60)
+    config = WaitForCallbackConfig(timeout_seconds=300, heartbeat_timeout_seconds=60)
     callback_id = create_callback_handler(
         state=mock_state,
         operation_identifier=OperationIdentifier("lifecycle_callback", None),
@@ -742,7 +742,7 @@ def mock_submitter(cb_id):
         assert cb_id == "lifecycle_cb123"
         return "submitted"
 
-    def execute_step(func, name):
+    def execute_step(func, name, config=None):
         step_context = Mock(spec=StepContext)
         return func(step_context)