feat: update wait_for_callback examples for new submitter signature

examples/src/wait_for_callback/wait_for_callback.py

@@ -1,12 +1,14 @@
 from typing import Any
 
-from aws_durable_execution_sdk_python.config import WaitForCallbackConfig
-from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    WaitForCallbackContext,
+)
 from aws_durable_execution_sdk_python.execution import durable_execution
-from aws_durable_execution_sdk_python.config import Duration
 
 
-def external_system_call(_callback_id: str) -> None:
+def external_system_call(_callback_id: str, _context: WaitForCallbackContext) -> None:
     """Simulate calling an external system with callback ID."""
     # In real usage, this would make an API call to an external system
     # passing the callback_id for the system to call back when done

examples/src/wait_for_callback/wait_for_callback_anonymous.py

@@ -10,7 +10,9 @@
 @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))
+    result: str = context.wait_for_callback(
+        lambda _callback_id, _context: time.sleep(1)
+    )
 
     return {
         "callbackResult": result,

examples/src/wait_for_callback/wait_for_callback_child.py

@@ -2,12 +2,12 @@
 
 from typing import Any
 
+from aws_durable_execution_sdk_python.config import Duration
 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
@@ -16,7 +16,7 @@ def child_context_with_callback(child_context: DurableContext) -> dict[str, Any]
     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"
+        lambda _callback_id, _context: None, name="child-callback-op"
     )
 
     return {
@@ -29,7 +29,7 @@ def child_context_with_callback(child_context: DurableContext) -> dict[str, Any]
 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"
+        lambda _callback_id, _context: None, name="parent-callback-op"
     )
 
     child_context_result: dict[str, Any] = context.run_in_child_context(

examples/src/wait_for_callback/wait_for_callback_heartbeat.py

@@ -3,13 +3,15 @@
 import time
 from typing import Any
 
-from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    WaitForCallbackContext,
+)
 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:
+def submitter(_callback_id: str, _context: WaitForCallbackContext) -> None:
     """Simulate long-running submitter function."""
     time.sleep(5)
     return None

examples/src/wait_for_callback/wait_for_callback_mixed_ops.py

@@ -3,9 +3,9 @@
 import time
 from typing import Any
 
+from aws_durable_execution_sdk_python.config import Duration
 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
@@ -19,7 +19,7 @@ def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
         name="fetch-user-data",
     )
 
-    def submitter(_) -> None:
+    def submitter(_callback_id, _context) -> None:
         """Submitter uses data from previous step."""
         time.sleep(0.1)
         return None

examples/src/wait_for_callback/wait_for_callback_multiple_invocations.py

@@ -2,9 +2,9 @@
 
 from typing import Any
 
+from aws_durable_execution_sdk_python.config import Duration
 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
@@ -14,7 +14,7 @@ def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
     context.wait(Duration.from_seconds(1), name="wait-invocation-1")
 
     # First callback operation
-    def first_submitter(callback_id: str) -> None:
+    def first_submitter(callback_id: str, _context) -> None:
         """Submitter for first callback."""
         print(f"First callback submitted with ID: {callback_id}")
         return None
@@ -34,7 +34,7 @@ def first_submitter(callback_id: str) -> None:
     context.wait(Duration.from_seconds(1), name="wait-invocation-2")
 
     # Second callback operation
-    def second_submitter(callback_id: str) -> None:
+    def second_submitter(callback_id: str, _context) -> None:
         """Submitter for second callback."""
         print(f"Second callback submitted with ID: {callback_id}")
         return None

examples/src/wait_for_callback/wait_for_callback_nested.py

@@ -2,12 +2,12 @@
 
 from typing import Any
 
+from aws_durable_execution_sdk_python.config import Duration
 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
@@ -16,7 +16,7 @@ def inner_child_context(inner_child_ctx: DurableContext) -> dict[str, Any]:
     inner_child_ctx.wait(Duration.from_seconds(5), name="deep-wait")
 
     nested_callback_result: str = inner_child_ctx.wait_for_callback(
-        lambda _: None,
+        lambda _callback_id, _context: None,
         name="nested-callback-op",
     )
 
@@ -30,7 +30,7 @@ def inner_child_context(inner_child_ctx: DurableContext) -> dict[str, Any]:
 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,
+        lambda _callback_id, _context: None,
         name="inner-callback-op",
     )
 
@@ -51,7 +51,7 @@ def outer_child_context(outer_child_ctx: DurableContext) -> dict[str, Any]:
 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,
+        lambda _callback_id, _context: None,
         name="outer-callback-op",
     )
 

examples/src/wait_for_callback/wait_for_callback_serdes.py

@@ -4,9 +4,9 @@
 from datetime import datetime
 from typing import Any, Optional, TypedDict
 
+from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
 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, WaitForCallbackConfig
 from aws_durable_execution_sdk_python.serdes import SerDes
 
 
@@ -75,7 +75,7 @@ def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
     )
 
     result: CustomData = context.wait_for_callback(
-        lambda _: None,
+        lambda _callback_id, _context: None,
         name="custom-serdes-callback",
         config=config,
     )

examples/src/wait_for_callback/wait_for_callback_submitter_failure.py

@@ -2,20 +2,20 @@
 
 from typing import Any
 
+from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
 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.retries import (
     RetryStrategyConfig,
     create_retry_strategy,
 )
-from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
 
 
 @durable_execution
 def handler(event: dict[str, Any], context: DurableContext) -> dict[str, Any]:
     """Handler demonstrating waitForCallback with submitter retry and exponential backoff."""
 
-    def submitter(callback_id: str) -> None:
+    def submitter(callback_id: str, _context) -> None:
         """Submitter function that can fail based on event parameter."""
         print(f"Submitting callback to external system - callbackId: {callback_id}")
         raise Exception("Simulated submitter failure")

examples/src/wait_for_callback/wait_for_callback_submitter_failure_catchable.py

@@ -3,20 +3,20 @@
 import time
 from typing import Any
 
+from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
 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.retries import (
     RetryStrategyConfig,
     create_retry_strategy,
 )
-from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
 
 
 @durable_execution
 def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
     """Handler demonstrating waitForCallback with failing submitter."""
 
-    def submitter(_) -> None:
+    def submitter(_callback_id, _context) -> None:
         """Submitter function that fails after a delay."""
         time.sleep(0.5)
         # Submitter fails

examples/src/wait_for_callback/wait_for_callback_timeout.py

@@ -2,10 +2,9 @@
 
 from typing import Any
 
+from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
 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
 
 
 @durable_execution
@@ -16,7 +15,7 @@ def handler(_event: Any, context: DurableContext) -> dict[str, Any]:
         timeout=Duration.from_seconds(1), heartbeat_timeout=Duration.from_seconds(2)
     )
 
-    def submitter(_) -> None:
+    def submitter(_callback_id, _context) -> None:
         """Submitter succeeds but callback never completes."""
         return None