examples: Add wait_for_callback examples

- Wait for callback examples
- Add a helper method in the runner to get all nested operations
- Update sam template

Author: Alex Wang

examples/examples-catalog.json

@@ -100,6 +100,105 @@
       },
       "path": "./src/wait_for_callback/wait_for_callback.py"
     },
+    {
+      "name": "Wait For Callback Success Anonymous",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_anonymous.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_anonymous.py"
+    },
+    {
+      "name": "Wait For Callback Heartbeat Sends",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_heartbeat.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_heartbeat.py"
+    },
+    {
+      "name": "Wait For Callback With Child Context",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_child.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_child.py"
+    },
+    {
+      "name": "Wait For Callback Mixed Ops",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_mixed_ops.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_mixed_ops.py"
+    },
+    {
+      "name": "Wait For Callback Multiple Invocations",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_multiple_invocations.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_multiple_invocations.py"
+    },
+    {
+      "name": "Wait For Callback Failing Submitter Catchable",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_submitter_failure_catchable.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_submitter_failure_catchable.py"
+    },
+    {
+      "name": "Wait For Callback Submitter Failure",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_submitter_failure.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_submitter_failure.py"
+    },
+    {
+      "name": "Wait For Callback Serdes",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_serdes.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_serdes.py"
+    },
+    {
+      "name": "Wait For Callback Nested",
+      "description": "Usage of context.wait_for_callback() to wait for external system responses",
+      "handler": "wait_for_callback_nested.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/wait_for_callback/wait_for_callback_nested.py"
+    },
     {
       "name": "Run in Child Context",
       "description": "Usage of context.run_in_child_context() to execute operations in isolated contexts",

examples/src/wait_for_callback/wait_for_callback_anonymous.py

@@ -0,0 +1,18 @@
+"""Demonstrates waitForCallback with anonymous (inline) submitter function."""
+
+import time
+from typing import Any
+
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import durable_execution
+
+
+@durable_execution
+def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
+    """Handler demonstrating waitForCallback with anonymous submitter."""
+    result: str = context.wait_for_callback(lambda _: time.sleep(1))
+
+    return {
+        "callbackResult": result,
+        "completed": True,
+    }

examples/src/wait_for_callback/wait_for_callback_child.py

@@ -0,0 +1,42 @@
+"""Demonstrates waitForCallback operations within child contexts."""
+
+from typing import Any
+
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    durable_with_child_context,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.config import Duration
+
+
+@durable_with_child_context
+def child_context_with_callback(child_context: DurableContext) -> dict[str, Any]:
+    """Child context containing wait and callback operations."""
+    child_context.wait(Duration.from_seconds(1), name="child-wait")
+
+    child_callback_result: str = child_context.wait_for_callback(
+        lambda _: None, name="child-callback-op"
+    )
+
+    return {
+        "childResult": child_callback_result,
+        "childProcessed": True,
+    }
+
+
+@durable_execution
+def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
+    """Handler demonstrating waitForCallback within child contexts."""
+    parent_result: str = context.wait_for_callback(
+        lambda _: None, name="parent-callback-op"
+    )
+
+    child_context_result: dict[str, Any] = context.run_in_child_context(
+        child_context_with_callback(), name="child-context-with-callback"
+    )
+
+    return {
+        "parentResult": parent_result,
+        "childContextResult": child_context_result,
+    }

examples/src/wait_for_callback/wait_for_callback_heartbeat.py

@@ -0,0 +1,31 @@
+"""Demonstrates sending heartbeats during long-running callback processing."""
+
+import time
+from typing import Any
+
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.config import Duration
+from aws_durable_execution_sdk_python.config import WaitForCallbackConfig
+
+
+def submitter(_callback_id: str) -> None:
+    """Simulate long-running submitter function."""
+    time.sleep(5)
+    return None
+
+
+@durable_execution
+def handler(event: dict[str, Any], context: DurableContext) -> dict[str, Any]:
+    """Handler demonstrating waitForCallback with heartbeat timeout."""
+
+    config = WaitForCallbackConfig(
+        timeout=Duration.from_seconds(120), heartbeat_timeout=Duration.from_seconds(15)
+    )
+
+    result: str = context.wait_for_callback(submitter, config=config)
+
+    return {
+        "callbackResult": result,
+        "completed": True,
+    }

examples/src/wait_for_callback/wait_for_callback_mixed_ops.py

@@ -0,0 +1,47 @@
+"""Demonstrates waitForCallback combined with steps, waits, and other operations."""
+
+import time
+from typing import Any
+
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.config import Duration
+
+
+@durable_execution
+def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
+    """Handler demonstrating waitForCallback mixed with other operations."""
+    # Mix waitForCallback with other operation types
+    context.wait(Duration.from_seconds(1), name="initial-wait")
+
+    step_result: dict[str, Any] = context.step(
+        lambda _: {"userId": 123, "name": "John Doe"},
+        name="fetch-user-data",
+    )
+
+    def submitter(_) -> None:
+        """Submitter uses data from previous step."""
+        time.sleep(0.1)
+        return None
+
+    callback_result: str = context.wait_for_callback(
+        submitter,
+        name="wait-for-callback",
+    )
+
+    context.wait(Duration.from_seconds(2), name="final-wait")
+
+    final_step: dict[str, Any] = context.step(
+        lambda _: {
+            "status": "completed",
+            "timestamp": int(time.time() * 1000),
+        },
+        name="finalize-processing",
+    )
+
+    return {
+        "stepResult": step_result,
+        "callbackResult": callback_result,
+        "finalStep": final_step,
+        "workflowCompleted": True,
+    }

examples/src/wait_for_callback/wait_for_callback_multiple_invocations.py

@@ -0,0 +1,53 @@
+"""Demonstrates multiple invocations tracking with waitForCallback operations across different invocations."""
+
+from typing import Any
+
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.config import Duration
+
+
+@durable_execution
+def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
+    """Handler demonstrating multiple invocations with waitForCallback operations."""
+    # First invocation - wait operation
+    context.wait(Duration.from_seconds(1), name="wait-invocation-1")
+
+    # First callback operation
+    def first_submitter(callback_id: str) -> None:
+        """Submitter for first callback."""
+        print(f"First callback submitted with ID: {callback_id}")
+        return None
+
+    callback_result_1: str = context.wait_for_callback(
+        first_submitter,
+        name="first-callback",
+    )
+
+    # Step operation between callbacks
+    step_result: dict[str, Any] = context.step(
+        lambda _: {"processed": True, "step": 1},
+        name="process-callback-data",
+    )
+
+    # Second invocation - another wait operation
+    context.wait(Duration.from_seconds(1), name="wait-invocation-2")
+
+    # Second callback operation
+    def second_submitter(callback_id: str) -> None:
+        """Submitter for second callback."""
+        print(f"Second callback submitted with ID: {callback_id}")
+        return None
+
+    callback_result_2: str = context.wait_for_callback(
+        second_submitter,
+        name="second-callback",
+    )
+
+    # Final invocation returns complete result
+    return {
+        "firstCallback": callback_result_1,
+        "secondCallback": callback_result_2,
+        "stepResult": step_result,
+        "invocationCount": "multiple",
+    }

examples/src/wait_for_callback/wait_for_callback_nested.py

@@ -0,0 +1,66 @@
+"""Demonstrates nested waitForCallback operations across multiple child context levels."""
+
+from typing import Any
+
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    durable_with_child_context,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.config import Duration
+
+
+@durable_with_child_context
+def inner_child_context(inner_child_ctx: DurableContext) -> dict[str, Any]:
+    """Inner child context with deep nested callback."""
+    inner_child_ctx.wait(Duration.from_seconds(5), name="deep-wait")
+
+    nested_callback_result: str = inner_child_ctx.wait_for_callback(
+        lambda _: None,
+        name="nested-callback-op",
+    )
+
+    return {
+        "nestedCallback": nested_callback_result,
+        "deepLevel": "inner-child",
+    }
+
+
+@durable_with_child_context
+def outer_child_context(outer_child_ctx: DurableContext) -> dict[str, Any]:
+    """Outer child context with inner callback and nested context."""
+    inner_result: str = outer_child_ctx.wait_for_callback(
+        lambda _: None,
+        name="inner-callback-op",
+    )
+
+    # Nested child context with another callback
+    deep_nested_result: dict[str, Any] = outer_child_ctx.run_in_child_context(
+        inner_child_context(),
+        name="inner-child-context",
+    )
+
+    return {
+        "innerCallback": inner_result,
+        "deepNested": deep_nested_result,
+        "level": "outer-child",
+    }
+
+
+@durable_execution
+def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
+    """Handler demonstrating nested waitForCallback operations across multiple levels."""
+    outer_result: str = context.wait_for_callback(
+        lambda _: None,
+        name="outer-callback-op",
+    )
+
+    nested_result: dict[str, Any] = context.run_in_child_context(
+        outer_child_context(),
+        name="outer-child-context",
+    )
+
+    return {
+        "outerCallback": outer_result,
+        "nestedResults": nested_result,
+    }