fix: acknowledge delivered invocation input

packages/aws-durable-execution-sdk-python-testing/docs/architecture.md

@@ -167,9 +167,9 @@ the handler would wait on a lane it needs in order to make progress.
 2. Invoke. The blocking call to the function runs off the lane, bounded
    by the invocation timeout. While it blocks, the handler's checkpoint
    and get state tasks run on the now free lane.
-3. Finish. A task on the lane records the response and releases the
-   gate. It moves `PENDING` to `PRE_INVOKE`, or a terminal result to
-   `COMPLETED`.
+3. Finish. A task on the lane validates the response, acknowledges its input
+   snapshot, and releases the gate. It moves `PENDING` to `PRE_INVOKE`, or a
+   terminal result to `COMPLETED`.
 
 ```mermaid
 sequenceDiagram
@@ -234,10 +234,17 @@ list in a checkpoint response. The delta for the next response is the
 operations whose `operation_last_touched_seq` is greater than
 `handler_seen_seq`.
 
-The watermark advances only when a checkpoint response is built. It never
-advances when invocation input is built. That is what makes retries
-correct. A failed invocation never moves the watermark. So the retry
-delivers the same state again.
+The watermark advances when a checkpoint response is built or a valid
+invocation response acknowledges its input snapshot, including a `PENDING`
+response that sent no checkpoint. Constructing invocation input does not
+acknowledge delivery: a failed delivery leaves its updates available for retry.
+
+Acknowledgment removes `UpdatedOperationIds` only when their last-touched
+sequence is covered by that snapshot. Newer versions, including another
+version of the same operation, remain pending. An invocation's acknowledgment
+never retreats a watermark already advanced by a checkpoint. Legacy update
+IDs without a stored version are treated as sequence zero and can be consumed
+by an equal-sequence acknowledgment.
 
 ```
       seq_counter    handler_seen_seq    delta next checkpoint

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py

@@ -479,15 +479,11 @@ def record_invocation_completion(
         )
 
     def mark_state_delivered(self) -> None:
-        """Reset the list of changed operations once an invocation's input
-        carries them.
-
-        The service reports, on each invocation, the operations that
-        changed after the last state the handler observed. The runner
-        delivers state in the invocation input. So the list is reset when
-        that input is built, not when the invocation completes: an
-        operation that completes while the handler is still running is
-        reported on the next invocation.
+        """Explicitly clear all pending operation IDs after their delivery.
+
+        The executor uses versioned ``advance_handler_seen`` acknowledgments
+        so failed invocations retain their input updates and changes newer
+        than the acknowledged snapshot remain pending.
         """
         self.updated_operation_ids = []
 
@@ -860,10 +856,9 @@ class OperationPaginatorState:
     * :meth:`unseen_operations` — strict ``> handler_seen_seq``
       delta filter, used by ``checkpoint_execution`` to decide which
       ops go back to the handler.
-    * :meth:`advance_handler_seen` — monotonic forward update of the
-      delivery watermark on the wrapped :class:`Execution`, called
-      after a checkpoint response is constructed and its ops are
-      known.
+    * :meth:`advance_handler_seen` — acknowledge delivered operation versions
+      without retreating the watermark, after constructing a checkpoint
+      response or validating an invocation's response.
 
 
     Pinning ``token_sequence`` and ``operations`` at construction
@@ -920,11 +915,18 @@ def unseen_operations(self) -> list[Operation]:
         ]
 
     def advance_handler_seen(self, seq: int) -> None:
-        """Advance the wrapped execution's ``handler_seen_seq`` to
-        ``seq`` if that represents forward progress. Monotonic:
-        smaller or equal values are ignored."""
+        """Acknowledge delivered updates without retreating the watermark.
+
+        Equal acknowledgments also consume legacy, unversioned update IDs.
+        Only versions covered by this delivery may be removed.
+        """
         if seq > self.execution.handler_seen_seq:
             self.execution.handler_seen_seq = seq
+        self.execution.updated_operation_ids = [
+            operation_id
+            for operation_id in self.execution.updated_operation_ids
+            if self.execution.operation_last_touched_seq.get(operation_id, 0) > seq
+        ]
 
     # --- internals -------------------------------------------------
 

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py

@@ -1551,7 +1551,8 @@ def _begin_invocation(
         self._set_invocation_gate(execution_arn, InvocationState.INVOKING)
         execution.begin_new_invocation()
         invocation_input = self._invoker.create_invocation_input(execution=execution)
-        execution.mark_state_delivered()
+        # Retain input updates until a valid response acknowledges this snapshot.
+        # A failed invocation must be able to receive the same updates on retry.
         self._store.save(execution)
         return execution, invocation_input, execution.seq_counter
 
@@ -1618,6 +1619,12 @@ def _finish_invocation(
         # A clean invocation resets the retry counter.
         reset_target = self._store.load(execution_arn)
         reset_target.consecutive_failed_invocation_attempts = 0
+        if invocation_seq is not None:
+            # A valid response acknowledges input even without a checkpoint.
+            # Preserve changes newer than the input and any newer checkpoint ack.
+            OperationPaginatorState.pin(reset_target).advance_handler_seen(
+                invocation_seq
+            )
         self._store.save(reset_target)
 
         # Release the gate: terminal -> COMPLETED, otherwise PRE_INVOKE

packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py

@@ -0,0 +1,183 @@
+"""A wait update survives failed delivery and a checkpoint-free resume."""
+
+import json
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from aws_durable_execution_sdk_python import (
+    DurableContext,
+    durable_execution,
+    durable_step,
+)
+from aws_durable_execution_sdk_python.config import CallbackConfig, Duration
+from aws_durable_execution_sdk_python.execution import (
+    DurableExecutionInvocationInput,
+    InvocationStatus,
+)
+from aws_durable_execution_sdk_python.plugin import (
+    DurableInstrumentationPlugin,
+    InvocationEndInfo,
+    InvocationStartInfo,
+    OperationChangeInfo,
+)
+from aws_durable_execution_sdk_python.serdes import JsonSerDes
+from aws_durable_execution_sdk_python_testing.execution import Execution
+from aws_durable_execution_sdk_python_testing.executor import Executor, InvocationState
+from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner
+from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore
+from aws_durable_execution_sdk_python_testing.stores.filesystem import (
+    FileSystemExecutionStore,
+)
+from aws_durable_execution_sdk_python_testing.stores.memory import (
+    InMemoryExecutionStore,
+)
+
+
+@pytest.mark.parametrize("filesystem", [False, True])
+@pytest.mark.parametrize("failed_delivery", [False, True])
+@pytest.mark.parametrize("legacy_format", [False, True])
+def test_wait_input_acknowledgment_survives_checkpoint_free_resume(
+    filesystem: bool,
+    failed_delivery: bool,
+    legacy_format: bool,
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    store: ExecutionStore = (
+        FileSystemExecutionStore.create(tmp_path / "store")
+        if filesystem
+        else InMemoryExecutionStore()
+    )
+    starts: list[InvocationStartInfo] = []
+    ends: list[InvocationEndInfo] = []
+    wait_deliveries: list[int] = []
+    attempts: list[int] = []
+    step_calls: list[str] = []
+    legacy_restores: list[dict[str, Any]] = []
+    original_begin = Executor._begin_invocation  # noqa: SLF001
+
+    def begin_with_legacy_state(
+        executor: Executor, arn: str
+    ) -> tuple[Execution, DurableExecutionInvocationInput, int] | None:
+        if (
+            legacy_format
+            and not legacy_restores
+            and executor._invocation_gate(arn) is InvocationState.PRE_INVOKE  # noqa: SLF001
+        ):
+            execution = store.load(arn)
+            if execution.updated_operation_ids:
+                # Rehydrate an actual wait completion in the older store format.
+                snapshot = execution.to_json_dict()
+                for field in (
+                    "SeqCounter",
+                    "HandlerSeenSeq",
+                    "OperationLastTouchedSeq",
+                    "OperationSizeBytes",
+                ):
+                    snapshot.pop(field)
+                restored = Execution.from_json_dict(snapshot)
+                assert restored.seq_counter == restored.handler_seen_seq == 0
+                assert restored.updated_operation_ids == execution.updated_operation_ids
+                store.save(restored)
+                legacy_restores.append(snapshot)
+        return original_begin(executor, arn)
+
+    monkeypatch.setattr(Executor, "_begin_invocation", begin_with_legacy_state)
+
+    class Observe(DurableInstrumentationPlugin):
+        def on_invocation_start(self, info: InvocationStartInfo) -> None:
+            starts.append(info)
+            if any(
+                op.name == "delivered-wait" and op.status.value == "SUCCEEDED"
+                for op in info.updated_operations.values()
+            ):
+                wait_deliveries.append(len(starts))
+
+        def on_invocation_end(self, info: InvocationEndInfo) -> None:
+            ends.append(info)
+
+        def on_operation_change(self, info: OperationChangeInfo) -> None:
+            if any(
+                op.name == "delivered-wait" and op.status.value == "SUCCEEDED"
+                for op in info.updated_operations.values()
+            ):
+                wait_deliveries.append(len(starts))
+
+    class ObserveFactory:
+        def create_plugin(
+            self, info: InvocationStartInfo
+        ) -> DurableInstrumentationPlugin:
+            return Observe()
+
+    @durable_step
+    def marker(_context: Any, value: str) -> str:
+        step_calls.append(value)
+        return value
+
+    def durable_handler(_event: Any, context: DurableContext) -> list[str]:
+        before = context.step(marker("before"), name="before-wait")
+        barrier = context.create_callback(
+            name="barrier", config=CallbackConfig(serdes=JsonSerDes())
+        )
+        context.wait(Duration.from_seconds(1), name="delivered-wait")
+        value = barrier.result()
+        assert isinstance(value, str)
+        return [before, context.step(marker(value), name="after-barrier")]
+
+    handler = durable_execution(durable_handler, plugins=[ObserveFactory()])
+
+    def entrypoint(event: Any, context: Any) -> Any:
+        attempts.append(1)
+        if failed_delivery and len(attempts) == 2:
+            raise RuntimeError("delivery failed before the SDK received its input")
+        return handler(event, context)
+
+    with DurableFunctionTestRunner(
+        handler=entrypoint,
+        store=store,
+        skip_time=False,
+        poll_interval=0.01,
+        execution_timeout=25,
+    ) as runner:
+        arn = runner.run_async(input="{}")
+        callback_id = runner.wait_for_callback(arn, name="barrier", timeout=5)
+        deadline = time.monotonic() + 10
+        while time.monotonic() < deadline:
+            events = runner.get_execution_history(arn).events
+            if (
+                len(ends) >= 2
+                and sum(event.event_type == "InvocationCompleted" for event in events)
+                >= 2
+            ):
+                break
+            time.sleep(0.01)
+        else:
+            raise AssertionError("The checkpoint-free resumed invocation did not end")
+        saved = store.load(arn)
+        wait = next(op for op in saved.operations if op.name == "delivered-wait")
+        acknowledged = (
+            list(saved.updated_operation_ids),
+            saved.handler_seen_seq,
+            saved.operation_last_touched_seq.get(wait.operation_id, 0),
+        )
+        runner.send_callback_success(callback_id, result=b'"after"')
+        result = runner.wait_for_result(arn, timeout=10)
+
+    assert result.status is InvocationStatus.SUCCEEDED
+    assert result.result is not None
+    assert json.loads(result.result) == ["before", "after"]
+    pending_ids, seen_sequence, wait_sequence = acknowledged
+    assert pending_ids == []
+    assert seen_sequence >= wait_sequence
+    assert [info.status.value for info in ends] == [
+        "PENDING",
+        "PENDING",
+        "SUCCEEDED",
+    ]
+    assert set(wait_deliveries) == {2}
+    assert step_calls == ["before", "after"]
+    assert len(attempts) == (4 if failed_delivery else 3)
+    assert len(legacy_restores) == int(legacy_format)

packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py

@@ -736,6 +736,9 @@ def test_pending_response_is_valid_when_an_operation_completed_after_the_handler
     mock_execution.consecutive_failed_invocation_attempts = 0
     mock_execution.seq_counter = 4
     mock_execution.handler_seen_seq = 2
+    mock_execution.operations = []
+    mock_execution.updated_operation_ids = []
+    mock_execution.operation_last_touched_seq = {}
     mock_execution.has_pending_operations.return_value = False
     mock_execution.has_changes_after.side_effect = lambda seq: seq == 4
     mock_execution.has_unseen_changes.return_value = True