feat: WaitForCallbackContext submitter with int tests

.github/workflows/ci.yml

@@ -42,7 +42,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install Hatch 
       run: |
-        python -m pip install --upgrade hatch
+        python -m pip install hatch==1.15.0
     - name: static analysis
       run: hatch fmt --check
     - name: type checking

.github/workflows/integration-tests.yml

@@ -2,7 +2,7 @@ name: Integration Tests
 
 on:
   pull_request:
-    branches: [ main ]
+    branches: [ main, runtime-deployed-version-2025-12-01 ]
 
 permissions:
   contents: read
@@ -42,7 +42,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
 
     - name: Install Hatch
-      run: python -m pip install --upgrade hatch
+      run: python -m pip install hatch==1.15.0
 
     - name: Setup and run Testing SDK
       working-directory: testing-sdk
@@ -107,7 +107,7 @@ jobs:
       env:
         AWS_DURABLE_SDK_URL: file://${{ github.workspace }}/language-sdk
       run: |
-        pip install hatch
+        python -m pip install hatch==1.15.0
         python -m pip install -e .
 
     - name: Get integration examples

.github/workflows/sync-package.yml

@@ -27,7 +27,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install Hatch
         run: |
-          python -m pip install --upgrade hatch
+          python -m pip install hatch==1.15.0
       - name: Build distribution
         run: hatch build
       - name: configure aws credentials

src/aws_durable_execution_sdk_python/context.py

@@ -47,6 +47,7 @@
     BatchResult,
     LoggerInterface,
     StepContext,
+    WaitForCallbackContext,
     WaitForConditionCheckContext,
 )
 from aws_durable_execution_sdk_python.types import Callback as CallbackProtocol
@@ -489,7 +490,7 @@ def wait(self, duration: Duration, name: str | None = None) -> None:
 
     def wait_for_callback(
         self,
-        submitter: Callable[[str], None],
+        submitter: Callable[[str, WaitForCallbackContext], None],
         name: str | None = None,
         config: WaitForCallbackConfig | None = None,
     ) -> Any:

src/aws_durable_execution_sdk_python/execution.py

@@ -97,7 +97,6 @@ class DurableExecutionInvocationInput:
     durable_execution_arn: str
     checkpoint_token: str
     initial_execution_state: InitialExecutionState
-    is_local_runner: bool
 
     @staticmethod
     def from_dict(
@@ -109,15 +108,13 @@ def from_dict(
             initial_execution_state=InitialExecutionState.from_dict(
                 input_dict.get("InitialExecutionState", {})
             ),
-            is_local_runner=input_dict.get("LocalRunner", False),
         )
 
     def to_dict(self) -> MutableMapping[str, Any]:
         return {
             "DurableExecutionArn": self.durable_execution_arn,
             "CheckpointToken": self.checkpoint_token,
             "InitialExecutionState": self.initial_execution_state.to_dict(),
-            "LocalRunner": self.is_local_runner,
         }
 
 
@@ -139,7 +136,6 @@ def from_durable_execution_invocation_input(
             durable_execution_arn=invocation_input.durable_execution_arn,
             checkpoint_token=invocation_input.checkpoint_token,
             initial_execution_state=invocation_input.initial_execution_state,
-            is_local_runner=invocation_input.is_local_runner,
             service_client=service_client,
         )
 
@@ -237,15 +233,12 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                 )
                 raise ExecutionError(msg) from e
 
-            # Local runner always uses its own client, otherwise use custom or default
-            if invocation_input.is_local_runner:
-                service_client = LambdaClient.initialize_local_runner_client()
-            else:
-                service_client = (
-                    LambdaClient(client=boto3_client)
-                    if boto3_client is not None
-                    else LambdaClient.initialize_from_env()
-                )
+            # Use custom client if provided, otherwise initialize from environment
+            service_client = (
+                LambdaClient(client=boto3_client)
+                if boto3_client is not None
+                else LambdaClient.initialize_from_env()
+            )
 
         raw_input_payload: str | None = (
             invocation_input.initial_execution_state.get_input_payload()

src/aws_durable_execution_sdk_python/lambda_service.py

@@ -951,33 +951,6 @@ def load_preview_botocore_models() -> None:
             Path(__file__).parent.joinpath("botocore", "data")
         )
 
-    @staticmethod
-    def initialize_local_runner_client() -> LambdaClient:
-        endpoint = os.getenv(
-            "DURABLE_LOCAL_RUNNER_ENDPOINT", "http://host.docker.internal:5000"
-        )
-        region = os.getenv("DURABLE_LOCAL_RUNNER_REGION", "us-west-2")
-
-        # The local runner client needs execute-api as the signing service name,
-        # so we have a second `lambdainternal-local` boto model with this.
-        LambdaClient.load_preview_botocore_models()
-        client = boto3.client(
-            "lambdainternal-local",
-            endpoint_url=endpoint,
-            region_name=region,
-            config=Config(
-                connect_timeout=5,
-                read_timeout=50,
-            ),
-        )
-
-        logger.debug(
-            "Initialized lambda client with endpoint: '%s', region: '%s'",
-            endpoint,
-            region,
-        )
-        return LambdaClient(client=client)
-
     @staticmethod
     def initialize_from_env() -> LambdaClient:
         LambdaClient.load_preview_botocore_models()

src/aws_durable_execution_sdk_python/operation/callback.py

@@ -10,6 +10,7 @@
     CallbackOptions,
     OperationUpdate,
 )
+from aws_durable_execution_sdk_python.types import WaitForCallbackContext
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -23,7 +24,11 @@
         CheckpointedResult,
         ExecutionState,
     )
-    from aws_durable_execution_sdk_python.types import Callback, DurableContext
+    from aws_durable_execution_sdk_python.types import (
+        Callback,
+        DurableContext,
+        StepContext,
+    )
 
 
 def create_callback_handler(
@@ -85,7 +90,7 @@ def create_callback_handler(
 
 def wait_for_callback_handler(
     context: DurableContext,
-    submitter: Callable[[str], None],
+    submitter: Callable[[str, WaitForCallbackContext], None],
     name: str | None = None,
     config: WaitForCallbackConfig | None = None,
 ) -> Any:
@@ -98,8 +103,10 @@ def wait_for_callback_handler(
         name=f"{name_with_space}create callback id", config=config
     )
 
-    def submitter_step(step_context):  # noqa: ARG001
-        return submitter(callback.callback_id)
+    def submitter_step(step_context: StepContext):
+        return submitter(
+            callback.callback_id, WaitForCallbackContext(logger=step_context.logger)
+        )
 
     step_config = (
         StepConfig(

src/aws_durable_execution_sdk_python/types.py

@@ -57,6 +57,11 @@ class StepContext(OperationContext):
     pass
 
 
+@dataclass(frozen=True)
+class WaitForCallbackContext(OperationContext):
+    """Context provided to waitForCallback submitter functions."""
+
+
 @dataclass(frozen=True)
 class WaitForConditionCheckContext(OperationContext):
     pass

tests/e2e/execution_int_test.py

@@ -68,7 +68,7 @@ def my_handler(event, context: DurableContext) -> list[str]:
         "aws_durable_execution_sdk_python.execution.LambdaClient"
     ) as mock_client_class:
         mock_client = Mock()
-        mock_client_class.initialize_local_runner_client.return_value = mock_client
+        mock_client_class.initialize_from_env.return_value = mock_client
 
         # Mock the checkpoint method to track calls
         checkpoint_calls = []
@@ -154,7 +154,7 @@ def my_handler(event, context: DurableContext):
         "aws_durable_execution_sdk_python.execution.LambdaClient"
     ) as mock_client_class:
         mock_client = Mock()
-        mock_client_class.initialize_local_runner_client.return_value = mock_client
+        mock_client_class.initialize_from_env.return_value = mock_client
 
         # Mock the checkpoint method to track calls
         checkpoint_calls = []
@@ -255,7 +255,7 @@ def my_handler(event, context):
         "aws_durable_execution_sdk_python.execution.LambdaClient"
     ) as mock_client_class:
         mock_client = Mock()
-        mock_client_class.initialize_local_runner_client.return_value = mock_client
+        mock_client_class.initialize_from_env.return_value = mock_client
 
         # Mock the checkpoint method to track calls
         checkpoint_calls = []
@@ -361,7 +361,7 @@ def my_handler(event, context: DurableContext):
         "aws_durable_execution_sdk_python.execution.LambdaClient"
     ) as mock_client_class:
         mock_client = Mock()
-        mock_client_class.initialize_local_runner_client.return_value = mock_client
+        mock_client_class.initialize_from_env.return_value = mock_client
 
         # Mock the checkpoint method to raise an error (using RuntimeError as a generic exception)
         def mock_checkpoint_failure(
@@ -424,7 +424,7 @@ def my_handler(event: Any, context: DurableContext):
         "aws_durable_execution_sdk_python.execution.LambdaClient"
     ) as mock_client_class:
         mock_client = Mock()
-        mock_client_class.initialize_local_runner_client.return_value = mock_client
+        mock_client_class.initialize_from_env.return_value = mock_client
 
         # Mock the checkpoint method to track calls
         checkpoint_calls = []