Set wait_for_callback context subtype

Author: kiwigitops

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py

@@ -717,6 +717,7 @@ def wait_in_child_context(context: DurableContext):
         return self.run_in_child_context(
             wait_in_child_context,
             step_name,
+            ChildConfig(sub_type=OperationSubType.WAIT_FOR_CALLBACK),
         )
 
     def wait_for_condition(

packages/aws-durable-execution-sdk-python/tests/context_test.py

@@ -1196,6 +1196,7 @@ def test_wait_for_callback_basic(mock_executor_class):
         # Verify the child context callable
         call_args = mock_run_in_child.call_args
         assert call_args[0][1] is None  # name should be None
+        assert call_args[0][2].sub_type is OperationSubType.WAIT_FOR_CALLBACK
 
 
 @patch("aws_durable_execution_sdk_python.context.wait_for_callback_handler")
@@ -1225,6 +1226,7 @@ def test_wait_for_callback_with_name_and_config(mock_executor_class):
         assert (
             call_args[0][1] == "submit_function"
         )  # name should be from _original_name
+        assert call_args[0][2].sub_type is OperationSubType.WAIT_FOR_CALLBACK
 
 
 @patch("aws_durable_execution_sdk_python.context.wait_for_callback_handler")
@@ -1250,6 +1252,7 @@ def test_wait_for_callback_resolves_name_from_submitter(mock_executor_class):
 
         call_args = mock_run_in_child.call_args
         assert call_args[0][1] == "submit_task"
+        assert call_args[0][2].sub_type is OperationSubType.WAIT_FOR_CALLBACK
 
 
 @patch("aws_durable_execution_sdk_python.context.wait_for_callback_handler")
@@ -1270,8 +1273,9 @@ def capture_handler_call(context, submitter, name, config):
 
     with patch.object(DurableContext, "run_in_child_context") as mock_run_in_child:
 
-        def run_child_context(callable_func, name):
+        def run_child_context(callable_func, name, config):
             # Execute the child context callable
+            assert config.sub_type is OperationSubType.WAIT_FOR_CALLBACK
             child_context = create_test_context(state=mock_state, parent_id="test")
             return callable_func(child_context)