fix: enforce cooperative invocation deadlines

.github/tests/otel_lifecycle_compatibility_test.py

@@ -125,6 +125,7 @@ def handler(event, durable_context):
             },
         }
         lambda_context = SimpleNamespace(
+            get_remaining_time_in_millis=lambda: 900000,
             aws_request_id="installed",
             client_context=None,
             identity=None,

packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py

@@ -61,6 +61,7 @@ def _execution_operation_id(owner: str) -> str:
 
 def _lambda_context(owner: str) -> Mock:
     context = Mock()
+    context.get_remaining_time_in_millis.return_value = 900000
     context.aws_request_id = f"request-{owner}"
     context.client_context = None
     context.identity = None

packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py

@@ -62,6 +62,7 @@
 
 def _lambda_context() -> Mock:
     context = Mock()
+    context.get_remaining_time_in_millis.return_value = 900000
     context.aws_request_id = "test-request-id"
     context.client_context = None
     context.identity = None

packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invoke_checkpoint_propagation.py

@@ -191,6 +191,7 @@ def user_handler(_event: Any, durable: DurableContext) -> list[str]:
         },
     }
     lambda_context = Mock()
+    lambda_context.get_remaining_time_in_millis.return_value = 900000
     lambda_context.aws_request_id = "request"
     lambda_context.client_context = None
     lambda_context.identity = None

packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_replaced_sampler_invoke.py

@@ -127,6 +127,7 @@ def user_handler(_event: Any, durable: DurableContext) -> str:
         invoked_function_arn="parent:1",
         tenant_id=None,
     )
+    lambda_context.get_remaining_time_in_millis.return_value = 900000
     service = Mock(checkpoint=checkpoint)
     before = context.get_current()
     try:

packages/aws-durable-execution-sdk-python/README.md

@@ -193,6 +193,50 @@ instances. Upgrade both packages together; core 3 rejects OTel 1.x registrations
 and the removed provider API version 1 contract. The independent registration
 API remains unchanged.
 
+### Invocation deadlines and Lambda Managed Instances
+
+The SDK captures `LambdaContext.get_remaining_time_in_millis()` at invocation
+entry and uses one monotonic budget, reserving 100 ms before the platform
+deadline for cooperative cleanup. An already-expired invocation starts no user
+handler. Once the budget is exhausted, new user operations and checkpoint/state
+requests are rejected with `InvocationDeadlineExceeded`, a retryable
+`InvocationError`. Step retry strategies do not turn this interruption into a
+persisted step failure. Completed checkpoints retain their stored results and
+errors for replay.
+
+Synchronous checkpoint waits, scheduled refresh waits, checkpoint queue reads,
+and map/parallel coordinator waits share that budget. In particular, a losing
+branch parked in an SDK refresh wait after early parallel completion unwinds at
+the cutoff, allowing invocation cleanup to join it. A newly created descendant
+of a completed context is rejected before its body, START hook, or checkpoint
+runs, even if it did not exist when its ancestor completed.
+
+Saved checkpoint reads and `ReplayChildren` reconstruction, including virtual
+children, remain available when a concurrent branch resumes in the same
+invocation. Reconstructing a saved result does not permit new durable work
+under that completed context.
+
+This is cooperative interruption, **not a guarantee of bounded invocation
+shutdown or LMI worker reclamation**. Python cannot safely terminate arbitrary
+running threads. The SDK still joins invocation-owned workers before returning:
+user code or I/O that does not cooperate, an already-running synchronous service
+client call (including its retries), or a blocked plugin hook can delay those
+joins indefinitely. The shared service client is not reconfigured or closed to
+interrupt one invocation. Give application I/O explicit finite timeouts and keep
+plugin hooks bounded. A runtime-supported worker-retirement strategy is still
+required for non-cooperative work; [issue #741](https://github.com/aws/aws-durable-execution-sdk-python/issues/741)
+tracks that remaining work, with real-LMI coverage in [test PR #743](https://github.com/aws/aws-durable-execution-sdk-python/pull/743).
+
+An expired invocation does not report an unacknowledged result as `SUCCEEDED` or
+`PENDING`. If the service has already acknowledged this invocation's terminal
+execution success/failure checkpoint, that outcome remains authoritative while
+local delivery hooks finish. A terminal checkpoint already in flight at the
+cutoff is settled before reporting the invocation outcome, including an
+acknowledgment received during cleanup. At the cutoff, further response metadata
+pages are skipped for that accepted terminal checkpoint; inline and already-read
+operations are retained. Already-running side effects cannot be undone;
+interrupted work continues to require the documented retry/idempotency strategy.
+
 ### Chained-invoke propagation
 
 A new invoke START collects synchronous plugin metadata after resolving its

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

@@ -292,6 +292,7 @@ def execute(
         execution_state.register_branch_pool(pool)
 
         def submit(branch: Branch[CallableType, ResultType]) -> None:
+            execution_state.invocation_deadline.check()
             branch.start()
             if execution_state._plugin_executor._plugins:  # noqa: SLF001
                 # Every admission/resume gets its own Context. Branch bindings
@@ -323,6 +324,7 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
             # and the reason, so user code runs at most once per change.
             decision: CompletionDecision | None = None
             while True:
+                execution_state.invocation_deadline.check()
                 if needs_snapshot_rebuild:
                     items_snapshot = (
                         self._build_items_snapshot()
@@ -386,7 +388,9 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
                 if resume_times:
                     timeout = max(resume_times[0] - time.time(), 0)
                 try:
-                    event: BranchEvent[ResultType] = events.get(timeout=timeout)
+                    event: BranchEvent[ResultType] = events.get(
+                        timeout=execution_state.invocation_deadline.remaining(timeout)
+                    )
                 except queue.Empty:
                     # A timed resume came due while branches were running.
                     continue

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

@@ -0,0 +1,49 @@
+"""One monotonic time budget for cooperative invocation interruption."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.exceptions import InvocationDeadlineExceeded
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.threading import CompletionEvent
+    from aws_durable_execution_sdk_python.types import LambdaContext
+
+# Stop admitting work before the platform deadline to leave time for cooperative
+# cleanup. This is not a bound on arbitrary user code, clients, or plugin hooks.
+_CLEANUP_MARGIN_SECONDS = 0.1
+
+
+@dataclass(frozen=True)
+class InvocationDeadline:
+    """Internal invocation-local deadline; None keeps standalone state unbounded."""
+
+    expires_at: float | None = None
+
+    @classmethod
+    def from_lambda_context(cls, context: LambdaContext) -> InvocationDeadline:
+        # Capture once: wall-clock adjustments cannot extend this invocation.
+        started_at = time.monotonic()
+        remaining_seconds = context.get_remaining_time_in_millis() / 1000
+        return cls(started_at + max(0.0, remaining_seconds - _CLEANUP_MARGIN_SECONDS))
+
+    def remaining(self, timeout: float | None = None) -> float | None:
+        """Clamp an existing wait to the same budget, without resetting it."""
+        if self.expires_at is None:
+            return timeout
+        remaining = max(0.0, self.expires_at - time.monotonic())
+        return remaining if timeout is None else min(timeout, remaining)
+
+    def check(self) -> None:
+        if self.remaining() == 0:
+            raise InvocationDeadlineExceeded("Lambda invocation time budget exhausted")
+
+    def wait(self, event: CompletionEvent, timeout: float | None = None) -> bool:
+        """Do not turn an acknowledged checkpoint into an error after its wait."""
+        completed = event.wait(self.remaining(timeout))
+        if not completed:
+            self.check()
+        return completed

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

@@ -104,6 +104,10 @@ def is_retryable(self) -> bool:
         return True
 
 
+class InvocationDeadlineExceeded(InvocationError):
+    """The invocation must retry after its cooperative time budget expires."""
+
+
 class DurableApiErrorCategory(Enum):
     INVOCATION = "INVOCATION"
     EXECUTION = "EXECUTION"

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

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING, Any
 
 from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.deadline import InvocationDeadline
 from aws_durable_execution_sdk_python.exceptions import (
     BackgroundThreadError,
     BotoClientError,
@@ -196,6 +197,8 @@ def durable_execution(
 
     @functools.wraps(func)
     def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
+        deadline = InvocationDeadline.from_lambda_context(context)
+        deadline.check()
         # Only the factory list outlives an invocation. The worker below owns
         # this fresh executor and all plugin lifecycle hooks.
         plugin_executor = PluginExecutor(configured_plugins)
@@ -234,6 +237,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
             service_client=service_client,
             plugin_executor=plugin_executor,
             updated_operation_ids=invocation_input.updated_operation_ids,
+            invocation_deadline=deadline,
         )
 
         try:
@@ -294,6 +298,7 @@ def invoke_handler() -> MutableMapping[str, Any]:
             )
             try:
                 # Background checkpointing errors will propagate through CompletionEvent.wait() as BackgroundThreadError
+                deadline.check()
                 result = func(input_event, durable_context)
 
                 # done with userland
@@ -496,6 +501,7 @@ def wait_for_checkpoint() -> None:
                         resources.callback(execution_state.close)
                         execution_operation = execution_state.get_execution_operation()
 
+                        deadline.check()
                         # execute the plugins
                         plugin_executor.on_invocation_start(
                             execution_arn=invocation_input.durable_execution_arn,
@@ -519,6 +525,7 @@ def wait_for_checkpoint() -> None:
                             execution_state.checkpoint_batches_forever
                         )
                         output = invoke_handler()
+                    execution_state.check_invocation_result()
                     # The ExitStack has now closed execution_state and joined the
                     # checkpoint thread. So any revocation observed while draining
                     # in-flight checkpoint work is visible here. A terminal status
@@ -571,7 +578,11 @@ def wait_for_checkpoint() -> None:
             else:
                 # Preserve the original fresh-worker context without plugins.
                 invocation_future = executor.submit(invoke_worker)
-            return invocation_future.result()
+            result = invocation_future.result()
+        # Cleanup hooks and all owned workers have returned. A hook can consume
+        # the remaining budget; do not call End again when this check fails.
+        execution_state.check_invocation_result()
+        return result
 
     return wrapper
 

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

@@ -109,6 +109,7 @@ class OperationExecutor(ABC, Generic[T]):
 
     state: ExecutionState
     operation_identifier: OperationIdentifier
+    is_virtual: bool = False
 
     def _get_checkpoint_result(self) -> CheckpointedResult:
         """Return this operation's checkpoint after validating replay identity."""
@@ -179,6 +180,9 @@ def process(self) -> T:
             InvalidStateError: If the check result is in an invalid state
             May raise operation-specific errors from check_result_status() or execute()
         """
+        self.state.check_operation_active(
+            self.operation_identifier, allow_replay=True, is_virtual=self.is_virtual
+        )
         # Check 1: Entry (handles replay and existing checkpoints)
         result = self.check_result_status()
 

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

@@ -204,6 +204,7 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
                 self.operation_identifier,
                 checkpointed_result.is_replay_children(),
                 attempt=None if checkpointed_result.is_existent() else 1,
+                is_virtual=self.is_virtual,
             )
             raw_result: T = wrapped_user_func()
 

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

@@ -12,6 +12,7 @@
 from aws_durable_execution_sdk_python.exceptions import (
     ExecutionError,
     InvalidStateError,
+    InvocationDeadlineExceeded,
     RetryableSerDesError,
     SerDesError,
     StepError,
@@ -266,6 +267,10 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
                 self.operation_identifier.name,
             )
             return return_value
+        except InvocationDeadlineExceeded:
+            # Deadline interruption belongs to the invocation, not the step
+            # retry policy or persisted operation failure.
+            raise
         except RetryableSerDesError:
             # Transient serdes failure: fail the invocation for backend retry,
             # bypassing the step retry strategy. For AT_MOST_ONCE steps the