diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index d0ab1f68..d112ccdf 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -350,7 +350,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: invocation_span.set_attribute( "durable.invocation.status", info.status.value ) - if info.status is InvocationStatus.FAILED: + if info.status in (InvocationStatus.FAILED, InvocationStatus.RETRY): invocation_span.set_status( StatusCode.ERROR, info.error.message if info.error else "" ) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 58d22a68..3c09004b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -175,7 +175,7 @@ def test_invocation_span_records_subsequent_invocation(): ("invocation_status", "expected_span_status"), [ (InvocationStatus.PENDING, StatusCode.UNSET), - (InvocationStatus.RETRY, StatusCode.UNSET), + (InvocationStatus.RETRY, StatusCode.ERROR), (InvocationStatus.SUCCEEDED, StatusCode.OK), (InvocationStatus.FAILED, StatusCode.ERROR), ], @@ -184,7 +184,7 @@ def test_invocation_span_status_reflects_execution_status( invocation_status: InvocationStatus, expected_span_status: StatusCode, ): - """Only terminal invocation spans receive a success or failure status.""" + """Invocation spans reflect successful, failed, and retry outcomes.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 7b280d46..222037a3 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -14,6 +14,7 @@ BackgroundThreadError, BotoClientError, CheckpointError, + DurableExecutionsError, ExecutionError, InvocationError, SuspendExecution, @@ -424,6 +425,11 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: # all user-space errors go here logger.exception("Execution failed") + if not isinstance(e, DurableExecutionsError): + # Preserve the terminal execution response while identifying + # the uncaught handler error in invocation telemetry. + plugin_executor.override_invocation_status(InvocationStatus.RETRY) + result = DurableExecutionInvocationOutput( status=InvocationStatus.FAILED, error=ErrorObject.from_exception(e) ).to_dict() 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 5294f1c7..127683c7 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 @@ -265,6 +265,7 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._plugins = plugins or [] self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None + self._invocation_status_override: InvocationStatus | None = None @contextlib.contextmanager def run(self): @@ -277,6 +278,7 @@ def run(self): yield finally: self._invocation_status = None + self._invocation_status_override = None # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) @@ -331,8 +333,13 @@ def on_invocation_start( is_first_invocation=is_first_invocation, execution_start_time=execution_start_time, ) + self._invocation_status_override = None self.execute_plugins(self._invocation_status, sync=True) + def override_invocation_status(self, status: InvocationStatus) -> None: + """Override the status reported to plugins without changing the output.""" + self._invocation_status_override = status + def on_invocation_end( self, output: "DurableExecutionInvocationOutput", @@ -341,6 +348,14 @@ def on_invocation_end( # on_invocation_start not called, skip return + if self._invocation_status_override is not None: + output = DurableExecutionInvocationOutput( + status=self._invocation_status_override, + result=output.result, + error=output.error, + ) + self._invocation_status_override = None + invocation_end_info = ( InvocationEndInfo.from_durable_execution_invocation_output( self._invocation_status, output diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index a03fe7f8..6d75631e 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -18,6 +18,7 @@ ExecutionError, GetExecutionStateError, InvocationError, + StepError, SuspendExecution, ) from aws_durable_execution_sdk_python.execution import ( @@ -2934,8 +2935,8 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert "invocation_end:SUCCEEDED" in plugin.calls -def test_durable_execution_with_plugins_failure(): - """Test that plugins receive invocation end and execution end on user error.""" +def test_durable_execution_with_plugins_uncaught_error(): + """Test that plugins classify an uncaught user error as a retry.""" mock_client = Mock(spec=DurableServiceClient) mock_output = CheckpointOutput( checkpoint_token="new_token", # noqa: S106 @@ -2955,6 +2956,31 @@ def test_handler(event: Any, context: DurableContext) -> dict: _make_lambda_context(), ) + assert result["Status"] == InvocationStatus.FAILED.value + assert "invocation_start" in plugin.calls + assert "invocation_end:RETRY" in plugin.calls + + +def test_durable_execution_with_plugins_operation_failure(): + """Test that plugins classify a durable operation error as failed.""" + mock_client = Mock(spec=DurableServiceClient) + mock_output = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + mock_client.checkpoint.return_value = mock_output + + plugin = _RecordingPlugin() + + @durable_execution(plugins=[plugin]) + def test_handler(event: Any, context: DurableContext) -> dict: + raise StepError("step failed") + + result = test_handler( + _make_invocation_input(mock_client), + _make_lambda_context(), + ) + assert result["Status"] == InvocationStatus.FAILED.value assert "invocation_start" in plugin.calls assert "invocation_end:FAILED" in plugin.calls