fix(plugin): fire the end hook on every exit

The invocation wrapper fired on_invocation_end on the success path and on
except Exception, so an exit by BaseException skipped it. A handler that
surfaces an asyncio.CancelledError -- user code awaiting a cancelled task
-- leaves that way, and so does a KeyboardInterrupt or SystemExit. The
invocation then ended without the one hook a plugin has to finish on:
Insight never drained, so the records it held for that execution were
dropped, and OTel never ended the spans it had opened. Teardown still ran,
so nothing failed visibly.

Every exit now fires the hook and re-raises the exception unchanged.

Plugin-code containment widens with it. The factory and hook boundaries
caught Exception, so a factory or hook raising CancelledError failed an
execution it was only observing and stopped the remaining plugins from
running. Both now contain every BaseException except KeyboardInterrupt,
SystemExit and GeneratorExit, which are instructions to the calling thread
rather than reports of a plugin defect. CancelledError is not among them:
nothing cancels the invocation thread or the single-worker plugin pool, so
one arriving from plugin code came from the plugin's own asyncio use.

ErrorObject.from_exception is annotated BaseException, which is what its
own helper already accepted.

Author: ParidelPooya

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

@@ -235,7 +235,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ErrorObject:
         )
 
     @classmethod
-    def from_exception(cls, exception: Exception) -> ErrorObject:
+    def from_exception(cls, exception: BaseException) -> ErrorObject:
         # SerDesError and subclasses pin to the base discriminator so replay
         # always reconstructs them as SerDesError.
         if isinstance(exception, SerDesError):

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

@@ -530,6 +530,21 @@ def _factory_name(factory: object) -> str:
     return getattr(factory, "__qualname__", None) or type(factory).__name__
 
 
+# Raised out of plugin code, these three are not reports of a plugin defect but
+# instructions to the thread that is running: stop. Containing one would drop the
+# instruction and return a thread that was told to unwind to the work after the
+# plugin. They are re-raised; every other BaseException is contained.
+#
+# asyncio.CancelledError is deliberately NOT here. It derives from BaseException
+# and it does mean "stop" for the task that was cancelled, but the task here is
+# the SDK's, not the plugin's: nothing cancels the invocation thread or the
+# single-worker plugin pool. A CancelledError arriving from plugin code therefore
+# came from the plugin's own asyncio use -- an awaited task it let be cancelled --
+# which is a plugin defect and belongs on the contained side, or the plugin's
+# failure would fail an execution it was only observing.
+_PLUGIN_THREAD_CONTROL_EXCEPTIONS = (KeyboardInterrupt, SystemExit, GeneratorExit)
+
+
 class PluginExecutor:
     """One invocation's plugin instances, metadata and dispatch.
 
@@ -611,12 +626,21 @@ def _create_plugins(self, info: InvocationStartInfo) -> None:
         :func:`plugin_discovery.load_configured_plugins` rejects such an entry
         while the handler is being initialized, so the silent case is not
         reachable through ``durable_execution()``.
+
+        Containment covers every ``BaseException`` except the three that instruct
+        the calling thread to stop; see
+        :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS`. Narrowing it to ``Exception``
+        left the contract conditional on a factory never raising outside that
+        hierarchy, and a factory that awaits a cancelled task raises
+        ``asyncio.CancelledError``, which is outside it.
         """
         plugins: list[DurableInstrumentationPlugin] = []
         for factory in self._plugin_factories:
             try:
                 plugin = factory.create_plugin(info)
-            except Exception:
+            except _PLUGIN_THREAD_CONTROL_EXCEPTIONS:
+                raise
+            except BaseException:  # noqa: BLE001 - a factory must not fail the execution
                 # log and ignore the exception
                 logger.exception(
                     "Plugin factory %s exception ignored", _factory_name(factory)
@@ -647,7 +671,14 @@ def _create_plugins(self, info: InvocationStartInfo) -> None:
 
     @staticmethod
     def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None:
-        """Invoke the appropriate plugin callback. Runs inside the thread pool."""
+        """Invoke the appropriate plugin callback. Runs inside the thread pool.
+
+        Contains every ``BaseException`` except the three that instruct the
+        calling thread to stop, the same rule the factory boundary uses. The
+        thread here is the executor's own single worker, which nothing outside
+        this class cancels or interrupts, so an exception outside the ``Exception``
+        hierarchy arriving here was raised by the plugin.
+        """
         try:
             match info:
                 case InvocationStartInfo():
@@ -666,7 +697,9 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None:
                     plugin.on_user_function_end(info)
                 case _:
                     raise RuntimeError(f"Unknown info type: {type(info)}")
-        except Exception:
+        except _PLUGIN_THREAD_CONTROL_EXCEPTIONS:
+            raise
+        except BaseException:  # noqa: BLE001 - a hook must not fail the execution
             # log and ignore the exception
             logger.exception("Plugin %s exception ignored", plugin.__class__.__name__)
 
@@ -1068,7 +1101,25 @@ def wrapper(event: Any, context: LambdaContext):
                             output=DurableExecutionInvocationOutput.from_dict(output),
                         )
                         return output
-                    except Exception as e:
+                    except BaseException as e:
+                        # Every exit fires the end hook, not only the ones that
+                        # derive from Exception. A handler that surfaces an
+                        # asyncio.CancelledError -- user code that awaited a
+                        # cancelled task, most simply -- leaves the invocation by
+                        # a BaseException, and an invocation that ends without
+                        # its end hook costs the plugins the only point at which
+                        # they can finish: Insight never drains, so the records it
+                        # holds for this execution are dropped, and OTel never
+                        # ends the spans it opened, so they are never exported.
+                        # The teardown below still runs either way, which is why
+                        # the gap was silent rather than a leak.
+                        #
+                        # KeyboardInterrupt and SystemExit reach here too, and
+                        # they also fire the hook. The hook is what a plugin needs
+                        # to flush, and a process being torn down is when flushing
+                        # matters; the cost is the same bounded work any
+                        # invocation end does. The exception itself is re-raised
+                        # unchanged, so what the caller sees is untouched.
                         plugin_executor.on_invocation_end(
                             output=DurableExecutionInvocationOutput.create_retry(
                                 ErrorObject.from_exception(e)

packages/aws-durable-execution-sdk-python/tests/plugin_test.py

@@ -1,3 +1,4 @@
+import asyncio
 import contextlib
 import datetime
 import logging
@@ -646,6 +647,41 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin:
         self.assertEqual(built[0].calls, ["invocation_start:req-1"])
         self.assertEqual(built[1].calls, ["invocation_start:req-2"])
 
+    def test_a_body_raising_cancellation_still_fires_the_end_hook(self):
+        """Every exit fires the end hook, not only the ones deriving from Exception.
+
+        A handler that surfaces an ``asyncio.CancelledError`` left the invocation
+        without its end hook, so Insight never drained the records it held for the
+        execution and OTel never ended the spans it had opened. Nothing failed
+        visibly, which is why the gap was silent.
+        """
+        for raised in (
+            asyncio.CancelledError("cancelled"),
+            KeyboardInterrupt(),
+            SystemExit(),
+        ):
+            with self.subTest(raised=type(raised).__name__):
+                plugin = _TrackingPlugin()
+                host = PluginHost(plugins=[plugin_factory(plugin)])
+
+                @host.handle_durable_output
+                def handler(event, context, plugin_executor):
+                    plugin_executor.on_invocation_start(
+                        execution_arn="arn:exec",
+                        lambda_context=LAMBDA_CTX,
+                        execution_start_time=START_TS,
+                        is_first_invocation=True,
+                    )
+                    raise raised
+
+                with self.assertRaises(type(raised)):
+                    handler({}, LAMBDA_CTX)
+
+                self.assertEqual(
+                    plugin.calls,
+                    ["invocation_start:req-1", "invocation_end:req-1"],
+                )
+
     def test_host_hands_out_a_new_executor_per_invocation(self):
         """The host itself holds no per-invocation state to overwrite."""
         host = PluginHost(plugins=[plugin_factory(_TrackingPlugin())])
@@ -891,6 +927,69 @@ def test_every_failing_factory_leaves_the_executor_usable(self):
                     ),
                 )
 
+    def test_a_factory_raising_cancellation_is_contained(self):
+        """Containment is not limited to ``Exception``.
+
+        ``asyncio.CancelledError`` derives from ``BaseException``, so a factory
+        that awaits a cancelled task used to abort the invocation it was only
+        instrumenting and stop the remaining factories from running.
+        """
+        surviving = _TrackingPlugin()
+
+        executor = PluginExecutor(
+            plugins=[_CancellingFactory(), plugin_factory(surviving)],
+        )
+
+        with self.assertLogs(
+            "aws_durable_execution_sdk_python.plugin", level=logging.ERROR
+        ) as logs:
+            with executor.run():
+                executor.on_invocation_start(
+                    execution_arn="arn:exec",
+                    lambda_context=LAMBDA_CTX,
+                    execution_start_time=START_TS,
+                    is_first_invocation=True,
+                )
+
+        self.assertIn("factory cancelled", "\n".join(logs.output))
+        self.assertEqual(surviving.calls, ["invocation_start:req-1"])
+
+    def test_a_factory_raising_thread_control_still_propagates(self):
+        """The three that tell the thread to stop are not contained.
+
+        Containing one would drop the instruction and hand the thread back to the
+        work that follows the plugin.
+        """
+        for control in (KeyboardInterrupt, SystemExit, GeneratorExit):
+            with self.subTest(control=control.__name__):
+                executor = PluginExecutor(
+                    plugins=[_ThreadControlFactory(control)],
+                )
+
+                with executor.run(), self.assertRaises(control):
+                    executor.on_invocation_start(
+                        execution_arn="arn:exec",
+                        lambda_context=LAMBDA_CTX,
+                        execution_start_time=START_TS,
+                        is_first_invocation=True,
+                    )
+
+    def test_a_hook_raising_cancellation_is_contained(self):
+        """The hook boundary uses the same rule as the factory boundary."""
+        tracking = _TrackingPlugin()
+        executor = PluginExecutor(
+            plugins=[plugin_factory(_CancellingPlugin()), plugin_factory(tracking)]
+        )
+
+        with self.assertLogs(
+            "aws_durable_execution_sdk_python.plugin", level=logging.ERROR
+        ) as logs:
+            with _invocation(executor, tracking):
+                executor.execute_plugins(OPERATION_START_INFO, sync=True)
+
+        self.assertIn("hook cancelled", "\n".join(logs.output))
+        self.assertIn("operation_start:op-2", tracking.calls)
+
 
 class TestPluginExecutor(unittest.TestCase):
     def test_no_thread_pool_when_plugins_is_none(self):
@@ -2171,6 +2270,30 @@ def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlug
         raise RuntimeError("factory boom")
 
 
+class _CancellingFactory:
+    """Factory whose ``create_plugin`` raises outside the ``Exception`` hierarchy."""
+
+    def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin:
+        raise asyncio.CancelledError("factory cancelled")
+
+
+class _ThreadControlFactory:
+    """Factory that raises one of the three exceptions that must propagate."""
+
+    def __init__(self, control: type[BaseException]) -> None:
+        self._control = control
+
+    def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin:
+        raise self._control
+
+
+class _CancellingPlugin(DurableInstrumentationPlugin):
+    """Plugin whose hook raises outside the ``Exception`` hierarchy."""
+
+    def on_operation_start(self, info):
+        raise asyncio.CancelledError("hook cancelled")
+
+
 class _FailingPlugin(DurableInstrumentationPlugin):
     """Plugin that raises on every hook call."""