fix: record callback terminal operation updates

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

@@ -692,6 +692,7 @@ def complete_callback_success(
                 end_timestamp=now if now is not None else real_now(),
                 callback_details=updated_callback_details,
             )
+            self._record_updated_operation(operation.operation_id)
             return self.operations[index]
 
     def complete_callback_failure(
@@ -718,6 +719,7 @@ def complete_callback_failure(
                 end_timestamp=now if now is not None else real_now(),
                 callback_details=updated_callback_details,
             )
+            self._record_updated_operation(operation.operation_id)
             return self.operations[index]
 
     def complete_callback_timeout(
@@ -744,6 +746,7 @@ def complete_callback_timeout(
                 end_timestamp=now if now is not None else real_now(),
                 callback_details=updated_callback_details,
             )
+            self._record_updated_operation(operation.operation_id)
             return self.operations[index]
 
     def complete_chained_invoke(

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

@@ -0,0 +1,128 @@
+"""Callback terminal transitions retain first-delivery update metadata."""
+
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Literal
+from unittest.mock import Mock
+
+import pytest
+from aws_durable_execution_sdk_python.lambda_service import (
+    CallbackDetails,
+    ErrorObject,
+    Operation,
+    OperationStatus,
+    OperationType,
+)
+
+from aws_durable_execution_sdk_python_testing.exceptions import IllegalStateException
+from aws_durable_execution_sdk_python_testing.execution import Execution
+from aws_durable_execution_sdk_python_testing.invoker import (
+    InProcessInvoker,
+    LambdaInvoker,
+)
+from aws_durable_execution_sdk_python_testing.model import StartDurableExecutionInput
+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,
+)
+
+
+Outcome = Literal["success", "failure", "timeout"]
+NOW = datetime(2026, 10, 27, tzinfo=UTC)
+ERROR = ErrorObject(
+    "callback failure", "CallbackRejected", "details", ["original frame"]
+)
+
+
+def _complete(execution: Execution, outcome: Outcome, callback_id: str) -> Operation:
+    if outcome == "success":
+        return execution.complete_callback_success(callback_id, b'"result"', now=NOW)
+    if outcome == "failure":
+        return execution.complete_callback_failure(callback_id, ERROR, now=NOW)
+    return execution.complete_callback_timeout(callback_id, ERROR, now=NOW)
+
+
+@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"])
+@pytest.mark.parametrize("filesystem", [False, True])
+@pytest.mark.parametrize("already_pending", [False, True])
+def test_callback_completion_records_only_successful_state_changes(
+    outcome: Outcome, filesystem: bool, already_pending: bool, tmp_path: Path
+) -> None:
+    target = Operation(
+        operation_id="callback",
+        operation_type=OperationType.CALLBACK,
+        status=OperationStatus.STARTED,
+        callback_details=CallbackDetails(callback_id="callback-id"),
+    )
+    other = Operation(
+        operation_id="other",
+        operation_type=OperationType.CALLBACK,
+        status=OperationStatus.SUCCEEDED,
+        callback_details=CallbackDetails(callback_id="other-id", result='"other"'),
+    )
+    execution = Execution(
+        "test-arn",
+        StartDurableExecutionInput(
+            account_id="123456789012",
+            function_name="test-function",
+            function_qualifier="$LATEST",
+            execution_name="test-execution",
+            execution_timeout_seconds=300,
+            execution_retention_period_days=7,
+            invocation_id="invocation-id",
+        ),
+        [target, other],
+    )
+    execution.updated_operation_ids = (
+        ["other", "callback"] if already_pending else ["other"]
+    )
+    before = execution.to_json_dict()
+    with pytest.raises(IllegalStateException):
+        _complete(execution, outcome, "missing")
+    assert execution.to_json_dict() == before
+
+    completed = _complete(execution, outcome, "callback-id")
+    assert execution.updated_operation_ids == ["other", "callback"]
+    assert execution.seq_counter == 1
+    assert execution.operation_last_touched_seq == {"callback": 1}
+    assert execution.token_sequence == 0
+    assert (
+        completed.status.value
+        == {"success": "SUCCEEDED", "failure": "FAILED", "timeout": "TIMED_OUT"}[
+            outcome
+        ]
+    )
+    assert completed.end_timestamp == NOW
+    assert completed.callback_details is not None
+    assert completed.callback_details.result == (
+        '"result"' if outcome == "success" else None
+    )
+    assert completed.callback_details.error == (None if outcome == "success" else ERROR)
+    assert execution.operations[1] == other
+
+    store: ExecutionStore = (
+        FileSystemExecutionStore.create(tmp_path)
+        if filesystem
+        else InMemoryExecutionStore()
+    )
+    store.save(execution)
+    restored = store.load(execution.durable_execution_arn)
+    assert restored.updated_operation_ids == ["other", "callback"]
+    assert restored.operations == execution.operations
+    assert restored.seq_counter == 1
+    assert restored.operation_last_touched_seq == {"callback": 1}
+    for invoker in [InProcessInvoker(Mock(), Mock()), LambdaInvoker(Mock())]:
+        invocation_input = invoker.create_invocation_input(restored)
+        assert invocation_input.updated_operation_ids == ["other", "callback"]
+        assert invocation_input.initial_execution_state.operations == [completed, other]
+
+    # Rejecting a second completion must not resurrect an already-delivered ID.
+    restored.mark_state_delivered()
+    before = restored.to_json_dict()
+    with pytest.raises(IllegalStateException, match="not in STARTED state"):
+        _complete(restored, outcome, "callback-id")
+    assert restored.to_json_dict() == before
+    assert restored.updated_operation_ids == []

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

@@ -0,0 +1,143 @@
+"""The first externally completed callback is visible before dependent code."""
+
+import json
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+from aws_durable_execution_sdk_python import DurableContext, durable_execution
+from aws_durable_execution_sdk_python.config import CallbackConfig, Duration
+from aws_durable_execution_sdk_python.exceptions import CallbackError
+from aws_durable_execution_sdk_python.lambda_service import ErrorObject
+from aws_durable_execution_sdk_python.plugin import (
+    DurableInstrumentationPlugin,
+    InvocationStartInfo,
+    OperationEndInfo,
+)
+from aws_durable_execution_sdk_python.serdes import JsonSerDes
+
+from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner
+from aws_durable_execution_sdk_python_testing.stores.filesystem import (
+    FileSystemExecutionStore,
+)
+
+
+def _suspended_callback(runner: DurableFunctionTestRunner, arn: str) -> str:
+    deadline = time.monotonic() + 10
+    while time.monotonic() < deadline:
+        events = runner.get_execution_history(arn, include_execution_data=True).events
+        starts = [
+            event
+            for event in events
+            if event.event_type == "CallbackStarted" and event.name == "target"
+        ]
+        if starts and any(
+            event.event_type == "InvocationCompleted"
+            and event.event_id > starts[0].event_id
+            for event in events
+        ):
+            details = starts[0].callback_started_details
+            assert details is not None and details.callback_id is not None
+            return details.callback_id
+        time.sleep(0.01)
+    raise AssertionError("The target callback did not suspend")
+
+
+@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"])
+@pytest.mark.parametrize("filesystem", [False, True])
+def test_callback_first_completion_precedes_dependent_step(
+    outcome: str, filesystem: bool, tmp_path: Path
+) -> None:
+    invocations: list[InvocationStartInfo] = []
+    ends: list[OperationEndInfo] = []
+    saved_calls: list[str] = []
+    observed_calls: list[str] = []
+    completions_before_step: list[OperationEndInfo] = []
+
+    class Observer(DurableInstrumentationPlugin):
+        def on_invocation_start(self, info: InvocationStartInfo) -> None:
+            invocations.append(info)
+
+        def on_operation_end(self, info: OperationEndInfo) -> None:
+            if info.name == "target":
+                ends.append(info)
+
+    class ObserverFactory:
+        def create_plugin(self, _info: InvocationStartInfo) -> Observer:
+            return Observer()
+
+    def saved_step(_context: Any) -> str:
+        saved_calls.append("saved")
+        return "saved"
+
+    def observed_step(_context: Any, value: str) -> str:
+        completions_before_step.extend(ends)
+        observed_calls.append(value)
+        return value
+
+    def body(_event: Any, context: DurableContext) -> str:
+        saved = context.step(saved_step, name="saved")
+        config = CallbackConfig(
+            timeout=Duration.from_seconds(1 if outcome == "timeout" else 30),
+            serdes=JsonSerDes(),
+        )
+        try:
+            value = context.create_callback(name="target", config=config).result()
+        except CallbackError:
+            value = outcome
+        assert isinstance(value, str)
+        observed = context.step(
+            lambda step_context: observed_step(step_context, value), name="observed"
+        )
+        return saved + "/" + observed
+
+    with DurableFunctionTestRunner(
+        handler=durable_execution(body, plugins=[ObserverFactory()]),
+        store=FileSystemExecutionStore.create(tmp_path) if filesystem else None,
+        skip_time=False,
+        poll_interval=0.01,
+        execution_timeout=20,
+    ) as runner:
+        arn = runner.run_async(input="{}")
+        callback_id = _suspended_callback(runner, arn)
+        if outcome == "success":
+            runner.send_callback_success(
+                callback_id, result=json.dumps("target").encode()
+            )
+        elif outcome == "failure":
+            runner.send_callback_failure(
+                callback_id, error=ErrorObject.from_message("explicit failure")
+            )
+        # Timeout exercises the scheduled public-runner completion path.
+        result = runner.wait_for_result(arn, timeout=10)
+
+    expected = "target" if outcome == "success" else outcome
+    assert result.status.value == "SUCCEEDED"
+    assert (
+        result.result is not None and json.loads(result.result) == "saved/" + expected
+    )
+    assert saved_calls == ["saved"]
+    assert observed_calls == [expected]
+    assert len(invocations) == 2
+    assert completions_before_step, "The terminal hook must precede dependent code"
+    terminal = completions_before_step[0]
+    assert terminal.is_replayed is False
+    assert (
+        terminal.status.value
+        == {"success": "SUCCEEDED", "failure": "FAILED", "timeout": "TIMED_OUT"}[
+            outcome
+        ]
+    )
+    snapshot = invocations[1].updated_operations[terminal.operation_id]
+    assert snapshot.status is terminal.status
+    assert snapshot.result == terminal.result
+    assert snapshot.error == terminal.error
+    if outcome == "success":
+        assert terminal.result == json.dumps("target")
+        assert terminal.error is None
+    elif outcome == "failure":
+        assert terminal.error is not None
+        assert terminal.error.message == "explicit failure"
+    else:
+        assert terminal.error is not None