fix(plugin): split exception groups, drain before the build

A BaseExceptionGroup is neither case the containment boundary tested. A group
carrying a KeyboardInterrupt is not an instance of one, so the tuple handler
naming the three did not match it and the broad handler contained the
interrupt inside it. Plugin code produces such a group without asking:
an asyncio.TaskGroup whose task is interrupted raises one.

Both boundaries now partition what plugin code raised. The control leaves are
returned to the caller and re-raised; what remains is logged as a contained
plugin failure. A group of ordinary failures is therefore contained whole,
and a mixed group is logged and then propagates only its control part. The
end-hook dispatch catches BaseException rather than naming the three, because
everything reaching it has already been partitioned.

Insight registers its invocation-end drain on entering the hook rather than
after building the record. _emit runs customer code, and a content transform
raising a BaseException escapes _apply_data_content, which contains only
Exception. That failure left the hook with the drain unrequested, so records
this execution had already scheduled stayed in a buffering exporter when the
environment froze.

Author: ParidelPooya

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

@@ -389,6 +389,16 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
             return
         emit_mode = self._shared._emit_mode
         with self._hook_frame():
+            # The drain is registered before anything can fail, not after the
+            # record is built. `_emit` runs customer code -- the content and
+            # result transforms, and `__del__` on an object a displaced record
+            # carried -- and a failure there leaves this hook by way of the SDK's
+            # containment. Registering afterwards meant such a failure skipped
+            # the drain, so records this execution had already scheduled stayed
+            # in a buffering exporter when the environment froze. Registering
+            # here costs nothing when the hook succeeds: the frame runs the drain
+            # once, on the way out, either way.
+            self._request_drain()
             with self._lock:
                 if not self._closed:
                     # Close the gate before emitting so a concurrent late hook for
@@ -447,12 +457,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
             # after this call is done, without waiting on records scheduled after the
             # call by other executions.
             #
-            # Asked for rather than performed here, so it runs when the outermost
-            # hook frame on this thread unwinds and every `_lock` hold is released.
-            # See `_hook_frame`: an invocation end that customer code re-entered
-            # from inside another hook's build would otherwise wait for the export
-            # worker while holding the lock that worker may need.
-            self._request_drain()
+            # Asked for rather than performed, so it runs when the outermost hook
+            # frame on this thread unwinds and every `_lock` hold is released. See
+            # `_hook_frame`: an invocation end that customer code re-entered from
+            # inside another hook's build would otherwise wait for the export
+            # worker while holding the lock that worker may need. The request
+            # itself is made at the top of this hook, so a failure in the build
+            # below cannot skip it.
 
     # -- emission -------------------------------------------------------------
 

packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py

@@ -19,12 +19,15 @@
 
 from __future__ import annotations
 
+import asyncio
 import datetime
 import itertools
 import threading
 import time
 from typing import Any
 
+import pytest
+
 from aws_durable_execution_sdk_python.lambda_service import (
     ErrorObject,
     OperationStatus,
@@ -680,6 +683,7 @@ class ConcurrentCaptureExporter:
 
     def __init__(self) -> None:
         self.records: list[dict[str, Any]] = []
+        self.flushes = 0
         self._lock = threading.Lock()
 
     def render(self, record: dict[str, Any]) -> Any:
@@ -690,7 +694,8 @@ def export(self, record: dict[str, Any]) -> None:
             self.records.append(record)
 
     def flush(self) -> None:
-        pass
+        with self._lock:
+            self.flushes += 1
 
     def snapshot(self) -> list[dict[str, Any]]:
         with self._lock:
@@ -1188,6 +1193,53 @@ def hook() -> None:
     )
 
 
+def test_an_end_transform_that_fails_still_drains_what_was_scheduled():
+    # `_emit` runs customer code between the snapshot and the hand-off: the content
+    # transforms, a result override, and __del__ on an object a displaced record
+    # carried. A failure there leaves the hook through the SDK's containment, and
+    # the drain used to be requested after the build, so that failure skipped it --
+    # leaving records this execution had already scheduled in a buffering exporter
+    # when the environment froze. The drain is now requested on entry.
+    exporter = ConcurrentCaptureExporter()
+    calls: list[str] = []
+
+    def failing_on_the_second_call(value: Any) -> Any:
+        calls.append("input")
+        if len(calls) > 1:
+            # CancelledError rather than an ordinary exception on purpose:
+            # _apply_data_content contains Exception so that a failing redactor
+            # cannot leak the raw value, and a BaseException is what escapes the
+            # build and leaves the hook.
+            raise asyncio.CancelledError("transform cancelled")
+        return value
+
+    factory = workflow_insight(
+        WorkflowInsightConfig(
+            exporters=[exporter],
+            emit_mode="on-change",
+            content=ContentConfig(input=failing_on_the_second_call),
+        )
+    )
+    start = _start(operations={})
+    plugin = factory.create_plugin(start)
+
+    # The first emit succeeds and schedules a RUNNING record.
+    plugin.on_invocation_start(start)
+
+    # The terminal emit fails inside the transform, so this hook raises. The SDK
+    # contains that; here it is raised directly, which is the same code path.
+    with pytest.raises(asyncio.CancelledError):
+        plugin.on_invocation_end(_end(operations=_ops(_step("s"))))
+
+    statuses = [record["status"] for record in exporter.snapshot()]
+    assert statuses == ["RUNNING"], (
+        "the record scheduled before the failing transform must have been drained "
+        f"to the exporters, not left pending: {statuses}"
+    )
+    assert exporter.flushes >= 1, "the drain must have flushed"
+    assert _wait_until(lambda: not factory._scheduler._worker_alive())
+
+
 def _free_for_another_thread(lock: Any) -> bool:
     """Report whether a lock is unheld, as seen from a thread that never took it.
 

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

@@ -545,6 +545,36 @@ def _factory_name(factory: object) -> str:
 _PLUGIN_THREAD_CONTROL_EXCEPTIONS = (KeyboardInterrupt, SystemExit, GeneratorExit)
 
 
+def _contain_plugin_failure(
+    error: BaseException, message: str, *message_args: object
+) -> BaseException | None:
+    """Log what plugin code raised, and return the part that must not be contained.
+
+    Returns ``None`` when the whole failure was contained, or the part that
+    instructs the calling thread to stop, which the caller re-raises.
+
+    A :class:`BaseExceptionGroup` is split rather than tested, because it is
+    neither of the two cases a plain ``isinstance`` chain covers: a group carrying
+    a :class:`KeyboardInterrupt` is not an instance of one, so a tuple handler
+    naming the three does not match it and a broad handler would contain the
+    interrupt inside it. Plugin code produces such a group without asking for it
+    -- an ``asyncio.TaskGroup`` whose task is interrupted raises one -- so the
+    group is partitioned: the control leaves are returned to be re-raised, and
+    what remains is logged like any other contained plugin failure.
+    """
+    control: BaseException | None
+    contained: BaseException | None
+    if isinstance(error, BaseExceptionGroup):
+        control, contained = error.split(_PLUGIN_THREAD_CONTROL_EXCEPTIONS)
+    elif isinstance(error, _PLUGIN_THREAD_CONTROL_EXCEPTIONS):
+        control, contained = error, None
+    else:
+        control, contained = None, error
+    if contained is not None:
+        logger.error(message, *message_args, exc_info=contained)
+    return control
+
+
 class PluginExecutor:
     """One invocation's plugin instances, metadata and dispatch.
 
@@ -632,24 +662,24 @@ def _create_plugins(self, info: InvocationStartInfo) -> None:
         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
+        Containment covers every ``BaseException`` except the parts 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
+        :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS` and
+        :func:`_contain_plugin_failure`. 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 _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)
+            except BaseException as error:  # noqa: BLE001 - a factory must not fail the execution
+                control = _contain_plugin_failure(
+                    error, "Plugin factory %s exception ignored", _factory_name(factory)
                 )
+                if control is not None:
+                    raise control from None
                 continue
             if plugin is None:
                 logger.error(
@@ -678,11 +708,11 @@ def _create_plugins(self, info: InvocationStartInfo) -> None:
     def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None:
         """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.
+        Contains every ``BaseException`` except the parts 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:
@@ -702,11 +732,12 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None:
                     plugin.on_user_function_end(info)
                 case _:
                     raise RuntimeError(f"Unknown info type: {type(info)}")
-        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__)
+        except BaseException as error:  # noqa: BLE001 - a hook must not fail the execution
+            control = _contain_plugin_failure(
+                error, "Plugin %s exception ignored", plugin.__class__.__name__
+            )
+            if control is not None:
+                raise control from None
 
     def execute_plugins(self, info, sync):
         """Dispatch one hook to this invocation's plugins.
@@ -725,14 +756,20 @@ def execute_plugins(self, info, sync):
         allocated what its end hook releases.
 
         The invocation-end hook is the one hook that finishes dispatching even
-        when a plugin raises one of those three. Every plugin it reaches has
-        already started, so cutting the loop short costs a plugin its only chance
-        to finish: Insight would not drain, and OTel would leave spans unended.
-        The first such exception is held and re-raised once every plugin has been
+        when a plugin raises one of those. Every plugin it reaches has already
+        started, so cutting the loop short costs a plugin its only chance to
+        finish: Insight would not drain, and OTel would leave spans unended. The
+        first such exception is held and re-raised once every plugin has been
         called, so the thread still stops and nothing is swallowed. No other hook
         defers: stopping a start-hook loop early leaves later plugins with nothing
         to clean up, because the pairing rule above then withholds their end hook
         too.
+
+        Anything :meth:`_dispatch_plugin` raises is already a thread-control
+        failure -- it contains everything else -- so the end path catches
+        ``BaseException`` rather than naming the three again. Naming them would
+        miss a :class:`BaseExceptionGroup` carrying one, which is what
+        :func:`_contain_plugin_failure` hands back.
         """
         if not self._executor:
             return
@@ -752,7 +789,7 @@ def execute_plugins(self, info, sync):
                 continue
             try:
                 self._dispatch_plugin(plugin, info)
-            except _PLUGIN_THREAD_CONTROL_EXCEPTIONS as control:
+            except BaseException as control:  # noqa: BLE001 - held and re-raised below
                 if deferred_control is None:
                     deferred_control = control
         if deferred_control is not None:

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

@@ -647,6 +647,64 @@ 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_factory_raising_a_group_with_a_control_exception_propagates(self):
+        """A group carrying a control exception is not contained.
+
+        ``BaseExceptionGroup`` is neither of the two cases an ``isinstance`` chain
+        covers: a group holding a ``KeyboardInterrupt`` is not an instance of one,
+        so naming the three in a handler does not match it and a broad handler
+        would swallow the interrupt inside it. Plugin code produces such a group
+        without asking for one -- an ``asyncio.TaskGroup`` whose task is
+        interrupted raises it.
+        """
+        group = BaseExceptionGroup(
+            "plugin group", [ValueError("contained"), KeyboardInterrupt()]
+        )
+        executor = PluginExecutor(plugins=[_GroupRaisingFactory(group)])
+
+        with (
+            self.assertLogs(
+                "aws_durable_execution_sdk_python.plugin", level=logging.ERROR
+            ) as logs,
+            executor.run(),
+            self.assertRaises(BaseExceptionGroup) as raised,
+        ):
+            executor.on_invocation_start(
+                execution_arn="arn:exec",
+                lambda_context=LAMBDA_CTX,
+                execution_start_time=START_TS,
+                is_first_invocation=True,
+            )
+
+        # Only the control leaf propagates; the rest was logged as a contained
+        # plugin failure.
+        self.assertEqual(
+            [type(leaf) for leaf in raised.exception.exceptions], [KeyboardInterrupt]
+        )
+        self.assertIn("contained", "\n".join(logs.output))
+
+    def test_a_factory_raising_a_group_without_a_control_exception_is_contained(self):
+        """A group of ordinary failures is contained like any other."""
+        surviving = _TrackingPlugin()
+        group = BaseExceptionGroup("plugin group", [ValueError("boom")])
+        executor = PluginExecutor(
+            plugins=[_GroupRaisingFactory(group), 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("boom", "\n".join(logs.output))
+        self.assertEqual(surviving.calls, ["invocation_start:req-1"])
+
     def test_an_end_hook_that_stops_the_thread_still_finishes_the_dispatch(self):
         """Every started plugin receives the end hook, then the thread stops.
 
@@ -2380,6 +2438,16 @@ def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlug
         raise self._control
 
 
+class _GroupRaisingFactory:
+    """Factory that raises a ``BaseExceptionGroup``, as an ``asyncio.TaskGroup`` does."""
+
+    def __init__(self, group: BaseExceptionGroup) -> None:
+        self._group = group
+
+    def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin:
+        raise self._group
+
+
 class _CancellingPlugin(DurableInstrumentationPlugin):
     """Plugin whose hook raises outside the ``Exception`` hierarchy."""