From 9452512ee9cd723d19096f69e91cd543f0c66307 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 23 Jul 2026 12:39:30 -0700 Subject: [PATCH] fix(plugin): order operation starts before attempts --- .../plugin.py | 10 +++-- .../aws_durable_execution_sdk_python/state.py | 15 +++---- .../tests/state_test.py | 41 ++++++++++++++++--- 3 files changed, 49 insertions(+), 17 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index d6c3aa4a..5294f1c7 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -212,8 +212,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def on_operation_start(self, info: OperationStartInfo) -> None: """ - Called when an operation checkpoints STARTED status, or when a prior - operation is replayed. This is called NOT within the thread that runs operation. + Called before an operation's START checkpoint is queued, or when a + prior operation is replayed. This guarantees that it strictly precedes + ``on_user_function_start``. This is called NOT within the thread that + runs operation. Args: info: Information about the operation. @@ -380,10 +382,10 @@ def on_operation_action( operation: Operation | None = None, previous_operation: Operation | None = None, ): - """Execute any registered plugins for a given operation when an update is checkpointed + """Execute registered plugins before an operation START is queued. Args: - update: the operation update that is checkpointed + update: The operation update being checkpointed. operation: the operation after the checkpoint previous_operation: the operation before the checkpoint """ diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index 299c1e86..baec099b 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -597,6 +597,14 @@ def create_checkpoint( CompletionEvent() if is_sync else None ) + if operation_update is not None: + # Dispatch before queueing so START strictly precedes any user + # function attempt, regardless of checkpoint synchronization mode. + self._plugin_executor.on_operation_action( + operation_update, + previous_operation=self.operations.get(operation_update.operation_id), + ) + # Create wrapper object for queue queued_op = QueuedOperation(operation_update, completion_event) @@ -759,13 +767,6 @@ def checkpoint_batches_forever(self) -> None: output.checkpoint_token, output.new_execution_state.next_marker, ) - for update in updates: - self._plugin_executor.on_operation_action( - update, - self.operations.get(update.operation_id), - previous_operations.get(update.operation_id), - ) - self._plugin_executor.on_operation_update( updated_operations, self.operations, diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py index 6a623f40..4fb277ff 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -3969,6 +3969,36 @@ def test_plugin_executor_on_operation_action_called_on_checkpoint(): assert plugin.operation_starts[0].is_replayed is False +def test_async_operation_start_precedes_user_function_start(): + """Async START notifies plugins before the user function begins.""" + plugin = _RecordingPlugin() + plugin_executor = PluginExecutor(plugins=[plugin]) + with plugin_executor.run(): + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=create_autospec(LambdaClient), + plugin_executor=plugin_executor, + ) + operation_id = "step-1" + operation_identifier = OperationIdentifier( + operation_id=operation_id, + sub_type=OperationSubType.STEP, + name="my-step", + ) + state.create_checkpoint( + OperationUpdate.create_step_start(operation_identifier), + is_sync=False, + ) + plugin_executor.on_user_function_start(operation_identifier, attempt=1) + + assert plugin.calls[:2] == [ + f"operation_start:{operation_id}", + f"user_function_start:{operation_id}", + ] + + def test_existing_operation_start_is_reported_as_replayed(): """A repeated START checkpoint reports that the operation already existed.""" mock_client = create_autospec(LambdaClient) @@ -4189,8 +4219,8 @@ def test_plugin_executor_called_for_multiple_updates_in_batch(): executor.shutdown(wait=True) # Both operations should have triggered on_operation_action - assert "operation_start:step-1" in plugin.calls - assert "operation_start:step-2" in plugin.calls + assert plugin.calls.count("operation_start:step-1") == 1 + assert plugin.calls.count("operation_start:step-2") == 1 # Both terminal operations should have triggered on_operation_update assert "operation_end:step-1" in plugin.calls assert "operation_end:step-2" in plugin.calls @@ -4264,8 +4294,8 @@ def test_plugin_executor_on_operation_change_called_for_status_changes(): assert "operation_change:wait-1" not in plugin.calls -def test_plugin_executor_not_called_on_checkpoint_failure(): - """Test that plugin_executor is NOT called when checkpoint API fails.""" +def test_operation_start_plugin_hook_fires_before_checkpoint_failure(): + """START is reported before queueing even when checkpointing later fails.""" mock_client = create_autospec(spec=LambdaClient) mock_client.checkpoint.side_effect = RuntimeError("API error") @@ -4298,8 +4328,7 @@ def test_plugin_executor_not_called_on_checkpoint_failure(): state.stop_checkpointing() executor.shutdown(wait=True) - # Plugin should NOT have been called since checkpoint failed - assert "operation_start:step-1" not in plugin.calls + assert plugin.calls.count("operation_start:step-1") == 1 assert "operation_end:step-1" not in plugin.calls