fix(insight): keep one pending export snapshot per execution

The export scheduler held a single pending slot for the whole plugin. When
two executions were in flight on one plugin instance, execution B's RUNNING
snapshot could replace execution A's SUCCEEDED snapshot before the worker
claimed it. A's drain() then returned once B's record was flushed, and A's
terminal record was never exported. Lambda runs one invocation per
environment, but the local test runner and the conformance suite drive
independent executions concurrently through one shared plugin, where this
lost 30-37% of terminal records in a two-thread probe.

Key the pending map by executionArn. A snapshot only replaces its own
execution's entry, the worker exports entries in first-arrival order, and a
flush request is honored only once every pending record is exported, so
drain() means everything scheduled by any execution has been delivered.
A RUNNING snapshot never replaces a pending terminal snapshot for the same
execution.

Close the execution state in on_invocation_end before emitting and
draining. An operation-change hook from a checkpoint completing during the
drain, or arriving after the state was discarded, is dropped instead of
recreating state and emitting a trailing RUNNING record.

When the worker thread cannot be started, drain() now exports pending
records inline on the calling thread instead of disabling export for the
plugin's lifetime and dropping the record.

Author: Alex Wang

packages/aws-durable-execution-sdk-python-insight/README.md

@@ -57,10 +57,14 @@ Behavior is validated cross-SDK by the `insight` conformance suite
 
 > **Note (asynchronous export).** Export rendering, truncation, `export()`, and
 > `flush()` run on one lazy background worker per plugin. Checkpoint hooks only
-> replace the latest pending snapshot and wake the worker. Consecutive
-> `on-change` snapshots may coalesce while an export is in flight. An invocation
-> that emits a record drains the latest snapshot and flushes exporters before it
-> returns; invocations that emit nothing do not start or flush the worker.
+> replace the latest pending snapshot **for their own execution** and wake the
+> worker; executions in flight at the same time (as under the local test runner)
+> never displace each other's snapshots, and a terminal snapshot is never
+> replaced by a later `RUNNING` one. Consecutive `on-change` snapshots of one
+> execution may coalesce while an export is in flight. An invocation that emits
+> a record drains every pending snapshot and flushes exporters before it
+> returns; invocations that emit nothing do not start or flush the worker. If
+> the worker thread cannot be started, that drain exports inline instead.
 
 ## Requirements
 

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

@@ -1,7 +1,17 @@
 # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
-"""Latest-pending asynchronous export scheduling for Workflow Insight."""
+"""Latest-pending asynchronous export scheduling for Workflow Insight.
+
+One pending slot is kept **per execution** (keyed by ``executionArn``). Each
+record is a complete snapshot of its execution, so a newer snapshot for the same
+execution supersedes an older one that has not been exported yet, while records
+for different executions never displace each other. The plugin already tracks
+state per execution, and the local test runner drives independent executions
+concurrently through one shared plugin instance, so a single plugin-wide slot
+would silently drop one execution's terminal record whenever another execution
+scheduled a snapshot first.
+"""
 
 from __future__ import annotations
 
@@ -15,70 +25,81 @@
 
 _logger = logging.getLogger("aws_durable_execution_sdk_python_insight")
 
+_TERMINAL_STATUSES = frozenset({"SUCCEEDED", "FAILED"})
+
+
+def _is_terminal(record: dict[str, Any]) -> bool:
+    return record.get("status") in _TERMINAL_STATUSES
+
 
 class _ExportScheduler:
-    """Run all exporters on one lazy worker with one latest pending record."""
+    """Run all exporters on one lazy worker with one pending record per execution."""
 
     def __init__(self, exporters: list[InsightExporter]) -> None:
         self._exporters = exporters
         self._condition = threading.Condition(threading.Lock())
-        self._pending: dict[str, Any] | None = None
+        # executionArn -> latest pending snapshot for that execution. Insertion
+        # ordered, so the worker exports executions in first-arrival order;
+        # replacing an entry keeps its position.
+        self._pending: dict[str, dict[str, Any]] = {}
         self._flush_requested = False
         self._flush_event: threading.Event | None = None
         self._worker: threading.Thread | None = None
-        self._disabled = False
+        self._start_failure_logged = False
 
     def schedule(self, record: dict[str, Any]) -> None:
-        """Replace the pending snapshot and return without running exporters."""
+        """Replace this execution's pending snapshot and return without exporting."""
+        key = str(record.get("executionArn", ""))
         displaced: dict[str, Any] | None = None
-        failed_pending: dict[str, Any] | None = None
         start_error: Exception | None = None
         with self._condition:
-            if self._disabled:
+            displaced = self._pending.get(key)
+            # A terminal snapshot is final. A RUNNING snapshot for the same
+            # execution that arrives after it (an operation-change hook from a
+            # checkpoint completing during the end-of-invocation drain) must not
+            # replace it, or the execution would be reported as still running.
+            if (
+                displaced is not None
+                and _is_terminal(displaced)
+                and not _is_terminal(record)
+            ):
                 return
-            displaced = self._pending
-            self._pending = record
-            failed_pending, start_error = self._ensure_worker_locked()
+            self._pending[key] = record
+            start_error = self._ensure_worker_locked()
             self._condition.notify()
-        # Releasing either record may run custom finalizers, so do it unlocked.
-        del displaced, failed_pending
-        if start_error is not None:
-            _logger.warning(
-                "workflow-insight: could not start export worker; disabling "
-                "asynchronous export: %s",
-                start_error,
-            )
+        # Releasing the displaced record may run custom finalizers, so do it unlocked.
+        del displaced
+        self._log_start_failure(start_error)
 
     def drain(self) -> None:
-        """Wait until the latest pending record is exported and exporters flush."""
-        failed_pending: dict[str, Any] | None = None
+        """Wait until every pending record is exported and exporters flush.
+
+        Records scheduled by any execution are exported before the flush, so a
+        drain issued at one execution's invocation end also delivers snapshots
+        that a concurrently running execution scheduled earlier.
+        """
         start_error: Exception | None = None
         with self._condition:
-            if self._disabled:
-                return
             if not self._flush_requested:
                 self._flush_requested = True
                 self._flush_event = threading.Event()
             flush_event = self._flush_event
             assert flush_event is not None
-            failed_pending, start_error = self._ensure_worker_locked()
-            started = not self._disabled
+            start_error = self._ensure_worker_locked()
+            worker_running = self._worker is not None
             self._condition.notify()
-        del failed_pending
-        if start_error is not None:
-            _logger.warning(
-                "workflow-insight: could not start export worker; disabling "
-                "asynchronous export: %s",
-                start_error,
-            )
-        if started:
+        self._log_start_failure(start_error)
+        if worker_running:
             flush_event.wait()
+            return
+        # No worker could be started. Export and flush on the calling thread so
+        # nothing scheduled is dropped; this is the invocation-end path, which
+        # already waits for delivery.
+        self._pump(flush_event)
 
-    def _ensure_worker_locked(
-        self,
-    ) -> tuple[dict[str, Any] | None, Exception | None]:
+    def _ensure_worker_locked(self) -> Exception | None:
         if self._worker is not None and self._worker.is_alive():
-            return None, None
+            return None
         worker = threading.Thread(
             target=self._run,
             name=f"workflow-insight-export-{id(self)}",
@@ -88,32 +109,45 @@ def _ensure_worker_locked(
         try:
             worker.start()
         except Exception as exc:  # noqa: BLE001 - instrumentation must not escape hooks
-            self._disabled = True
+            # Leave the pending records in place: drain() exports them inline,
+            # and a later schedule() retries starting a worker.
             self._worker = None
-            failed_pending = self._pending
-            self._pending = None
-            failed_event = self._flush_event
-            self._flush_event = None
-            self._flush_requested = False
-            if failed_event is not None:
-                failed_event.set()
-            return failed_pending, exc
-        return None, None
+            return exc
+        return None
+
+    def _log_start_failure(self, start_error: Exception | None) -> None:
+        if start_error is None or self._start_failure_logged:
+            return
+        self._start_failure_logged = True
+        _logger.warning(
+            "workflow-insight: could not start export worker; records are "
+            "exported inline at invocation end instead: %s",
+            start_error,
+        )
+
+    def _pop_pending_locked(self) -> dict[str, Any] | None:
+        if not self._pending:
+            return None
+        key = next(iter(self._pending))
+        return self._pending.pop(key)
+
+    def _take_flush_locked(self) -> threading.Event | None:
+        flush_event = self._flush_event
+        self._flush_event = None
+        self._flush_requested = False
+        return flush_event
 
     def _run(self) -> None:
         while True:
             record: dict[str, Any] | None = None
             flush_event: threading.Event | None = None
             with self._condition:
-                while self._pending is None and not self._flush_requested:
+                while not self._pending and not self._flush_requested:
                     self._condition.wait()
-                if self._pending is not None:
-                    record = self._pending
-                    self._pending = None
-                else:
-                    flush_event = self._flush_event
-                    self._flush_event = None
-                    self._flush_requested = False
+                record = self._pop_pending_locked()
+                if record is None:
+                    # Every pending record is exported: honor the flush request.
+                    flush_event = self._take_flush_locked()
 
             if record is not None:
                 self._export(record)
@@ -123,10 +157,25 @@ def _run(self) -> None:
             if flush_event is not None:
                 flush_event.set()
             with self._condition:
-                if self._pending is None and not self._flush_requested:
+                if not self._pending and not self._flush_requested:
                     self._worker = None
                     return
 
+    def _pump(self, flush_event: threading.Event) -> None:
+        """Export every pending record, then flush, on the calling thread."""
+        while True:
+            with self._condition:
+                record = self._pop_pending_locked()
+                if record is None:
+                    # Another inline drain may already have taken the request;
+                    # flushing twice is harmless, losing a record is not.
+                    self._take_flush_locked()
+            if record is None:
+                break
+            self._export(record)
+        self._flush()
+        flush_event.set()
+
     def _export(self, record: dict[str, Any]) -> None:
         for exporter in self._exporters:
             try:
@@ -159,4 +208,4 @@ def _worker_alive(self) -> bool:
 
     def _pending_count(self) -> int:
         with self._condition:
-            return int(self._pending is not None)
+            return len(self._pending)

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

@@ -160,7 +160,7 @@ def _apply_result_override(
 
 
 class _ExecutionState:
-    __slots__ = ("start_time", "parsed_arn", "cached_input", "operations")
+    __slots__ = ("start_time", "parsed_arn", "cached_input", "operations", "closed")
 
     def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None:
         self.start_time = start_time
@@ -169,6 +169,10 @@ def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None:
         # operation_id -> OperationInfo, adopted verbatim from the SDK's
         # authoritative snapshot (invocation start/end and operation-change).
         self.operations: dict[str, OperationInfo] = {}
+        # Set by on_invocation_end before it emits. An operation-change hook
+        # from a checkpoint that completes during the end-of-invocation drain
+        # is dropped, so no RUNNING snapshot can follow the final record.
+        self.closed = False
 
 
 class WorkflowInsightPlugin(DurableInstrumentationPlugin):
@@ -230,6 +234,23 @@ def _discard_state(self, execution_arn: str) -> None:
         with self._lock:
             self._state.pop(execution_arn, None)
 
+    def _open_state(self, execution_arn: str) -> _ExecutionState | None:
+        """Return the execution's state only while its invocation is open.
+
+        Unlike ``_ensure_state`` this never creates state: an operation-change
+        hook always follows an invocation start, so missing state means the
+        invocation already ended and its state was discarded.
+        """
+        with self._lock:
+            state = self._state.get(execution_arn)
+            if state is None or state.closed:
+                return None
+            return state
+
+    def _close_state(self, state: _ExecutionState) -> None:
+        with self._lock:
+            state.closed = True
+
     def _adopt_operations(
         self, state: _ExecutionState, operations: dict[str, OperationInfo]
     ) -> None:
@@ -270,7 +291,11 @@ def on_operation_change(self, info: OperationChangeInfo) -> None:
         arn = info.execution_arn
         if not arn or not self._sampled_in(arn):
             return
-        state = self._ensure_state(arn)
+        state = self._open_state(arn)
+        if state is None:
+            # The invocation already ended (or is draining its final record);
+            # a snapshot from a checkpoint that completed late is stale.
+            return
         # Replace state with the full operations snapshot carried by the hook.
         self._adopt_operations(state, info.operations)
         # on-change mode exports an updated RUNNING record on each change so
@@ -297,6 +322,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
         # Refresh from the fresh end-of-invocation snapshot before emitting so
         # the terminal record reflects the final operation map.
         self._adopt_operations(state, info.operations)
+        # Close before emitting: an operation-change hook arriving from a
+        # checkpoint that completes during the drain below is rejected, so no
+        # RUNNING snapshot can follow (or replace) the final record.
+        self._close_state(state)
         status = _STATUS_MAP.get(info.status, "RUNNING")
         is_terminal = status in ("SUCCEEDED", "FAILED")
         is_failure = status == "FAILED"