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

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 that execution's latest pending snapshot and wake the worker.
+> Consecutive `on-change` snapshots for one execution coalesce whenever the
+> worker has not yet claimed the previous one, including when no export is
+> running, so under fast in-memory checkpoints (the local runner) a burst may
+> deliver only its first and latest snapshots. Snapshots from different
+> executions never coalesce with each other. 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.
 
 ## Requirements
 

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

@@ -17,27 +17,38 @@
 
 
 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.
+
+    The pending map is keyed by ``executionArn`` so a snapshot only ever
+    replaces its own execution's pending snapshot. Several executions may be in
+    flight on one plugin instance (the local runner drives independent
+    executions concurrently through a shared ``workflow_insight(...)``), and a
+    single shared slot let execution B's ``RUNNING`` snapshot displace
+    execution A's terminal snapshot before the worker claimed it.
+    """
 
     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, in first-arrival order.
+        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
 
     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."""
         displaced: dict[str, Any] | None = None
-        failed_pending: dict[str, Any] | None = None
+        failed_pending: dict[str, dict[str, Any]] | None = None
         start_error: Exception | None = None
+        arn = str(record.get("executionArn", ""))
         with self._condition:
             if self._disabled:
                 return
-            displaced = self._pending
-            self._pending = record
+            displaced = self._pending.get(arn)
+            # Assigning an existing key keeps its first-arrival position.
+            self._pending[arn] = record
             failed_pending, start_error = self._ensure_worker_locked()
             self._condition.notify()
         # Releasing either record may run custom finalizers, so do it unlocked.
@@ -50,8 +61,8 @@ def schedule(self, record: dict[str, Any]) -> None:
             )
 
     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."""
+        failed_pending: dict[str, dict[str, Any]] | None = None
         start_error: Exception | None = None
         with self._condition:
             if self._disabled:
@@ -76,7 +87,7 @@ def drain(self) -> None:
 
     def _ensure_worker_locked(
         self,
-    ) -> tuple[dict[str, Any] | None, Exception | None]:
+    ) -> tuple[dict[str, dict[str, Any]] | None, Exception | None]:
         if self._worker is not None and self._worker.is_alive():
             return None, None
         worker = threading.Thread(
@@ -91,7 +102,7 @@ def _ensure_worker_locked(
             self._disabled = True
             self._worker = None
             failed_pending = self._pending
-            self._pending = None
+            self._pending = {}
             failed_event = self._flush_event
             self._flush_event = None
             self._flush_requested = False
@@ -105,11 +116,13 @@ def _run(self) -> None:
             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
+                if self._pending:
+                    # Export every pending execution before honoring a flush, so
+                    # drain() means everything scheduled so far was delivered.
+                    arn = next(iter(self._pending))
+                    record = self._pending.pop(arn)
                 else:
                     flush_event = self._flush_event
                     self._flush_event = None
@@ -123,7 +136,7 @@ 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
 
@@ -159,4 +172,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/tests/test_export_scheduler.py

@@ -13,8 +13,13 @@
 )
 
 
-def _record(value: str) -> dict[str, Any]:
-    return {"status": "RUNNING", "value": value, "operations": []}
+def _record(value: str, arn: str = "exec-a") -> dict[str, Any]:
+    return {
+        "executionArn": arn,
+        "status": "RUNNING",
+        "value": value,
+        "operations": [],
+    }
 
 
 def _wait_until(predicate, timeout: float = 5.0) -> bool:
@@ -157,3 +162,94 @@ def test_drain_waits_for_blocked_exporter() -> None:
 
     assert not drain_thread.is_alive()
     assert exporter.calls == [("export", "terminal"), ("flush", None)]
+
+
+# -- pending map is keyed per execution (#719 review, issue 1) ----------------
+
+
+def test_pending_is_keyed_per_execution() -> None:
+    exporter = BlockingExporter()
+    scheduler = _ExportScheduler([exporter])
+    scheduler.schedule(_record("a-first", arn="exec-a"))
+    assert exporter.started.wait(5.0)
+
+    # B's snapshot must not displace A's pending terminal snapshot.
+    scheduler.schedule(_record("a-terminal", arn="exec-a"))
+    scheduler.schedule(_record("b-running", arn="exec-b"))
+    assert scheduler._pending_count() == 2
+
+    exporter.release.set()
+    scheduler.drain()
+    assert exporter.calls == [
+        ("export", "a-first"),
+        ("export", "a-terminal"),
+        ("export", "b-running"),
+        ("flush", None),
+    ]
+
+
+def test_drain_delivers_every_pending_execution_before_flush() -> None:
+    exporter = BlockingExporter()
+    scheduler = _ExportScheduler([exporter])
+    scheduler.schedule(_record("a", arn="exec-a"))
+    assert exporter.started.wait(5.0)
+    scheduler.schedule(_record("b", arn="exec-b"))
+    scheduler.schedule(_record("c", arn="exec-c"))
+
+    exporter.release.set()
+    scheduler.drain()
+
+    exported = [value for kind, value in exporter.calls if kind == "export"]
+    assert exported == ["a", "b", "c"]
+    assert exporter.calls[-1] == ("flush", None)
+
+
+def test_concurrent_executions_each_deliver_terminal_record() -> None:
+    rounds = 50
+    exporter = CaptureExporter()
+    scheduler = _ExportScheduler([exporter])
+
+    def drive(arn: str) -> None:
+        for i in range(rounds):
+            scheduler.schedule(_record(f"{arn}-running-{i}", arn=arn))
+            scheduler.schedule(_record(f"{arn}-terminal-{i}", arn=arn))
+            scheduler.drain()
+
+    threads = [threading.Thread(target=drive, args=(arn,)) for arn in ("a", "b")]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(30.0)
+    assert not any(thread.is_alive() for thread in threads)
+
+    exported = [
+        value
+        for kind, value in exporter.calls
+        if kind == "export" and value is not None
+    ]
+    for arn in ("a", "b"):
+        terminals = [v for v in exported if v.startswith(f"{arn}-terminal-")]
+        assert terminals == [f"{arn}-terminal-{i}" for i in range(rounds)]
+
+
+# -- coalescing contract (#719 review, issue 2) -------------------------------
+
+
+def test_same_execution_coalesces_to_first_and_latest_while_blocked() -> None:
+    exporter = BlockingExporter()
+    scheduler = _ExportScheduler([exporter])
+    scheduler.schedule(_record("first"))
+    assert exporter.started.wait(5.0)
+
+    scheduler.schedule(_record("second"))
+    scheduler.schedule(_record("third"))
+    scheduler.schedule(_record("latest"))
+    assert scheduler._pending_count() == 1
+
+    exporter.release.set()
+    scheduler.drain()
+    assert exporter.calls == [
+        ("export", "first"),
+        ("export", "latest"),
+        ("flush", None),
+    ]

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

@@ -13,8 +13,10 @@
 from __future__ import annotations
 
 import datetime
+import threading
 from typing import Any
 
+import pytest
 from aws_durable_execution_sdk_python.lambda_service import (
     ErrorObject,
     OperationStatus,
@@ -368,10 +370,14 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin():
     assert rec["durationMs"] is not None and rec["durationMs"] >= 0
 
 
-# -- on-change schedules progress and delivers terminal state ----------------
+# -- on-change delivers every record when nothing coalesces ------------------
 
 
-def test_on_change_schedules_running_and_delivers_terminal():
+def test_on_change_delivers_every_record_when_drained_between_hooks():
+    # Draining after each hook removes the worker race, so this pins the
+    # "no back-pressure means no loss" contract: one RUNNING per start/change
+    # plus the terminal record, in order. (The coalescing contract itself is
+    # pinned in test_export_scheduler.py with a blocked exporter.)
     exporter = CaptureExporter()
     plugin = workflow_insight(
         WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change")
@@ -380,28 +386,73 @@ def test_on_change_schedules_running_and_delivers_terminal():
     op2 = _step("s2", op_id="2")
 
     plugin.on_invocation_start(_start(operations={}))
+    plugin._scheduler.drain()
     plugin.on_operation_change(
         OperationChangeInfo(
             execution_arn=ARN, updated_operations=_ops(op1), operations=_ops(op1)
         )
     )
+    plugin._scheduler.drain()
     plugin.on_operation_change(
         OperationChangeInfo(
             execution_arn=ARN, updated_operations=_ops(op2), operations=_ops(op1, op2)
         )
     )
+    plugin._scheduler.drain()
     plugin.on_invocation_end(_end(operations=_ops(op1, op2)))
 
     statuses = [record["status"] for record in exporter.records]
-    assert statuses
-    assert statuses[-1] == "SUCCEEDED"
-    assert set(statuses[:-1]) <= {"RUNNING"}
+    assert statuses == ["RUNNING", "RUNNING", "RUNNING", "SUCCEEDED"]
     final = exporter.records[-1]
     assert [op["name"] for op in final["operations"]] == ["s1", "s2"]
     ids = [op["id"] for op in final["operations"]]
     assert len(ids) == len(set(ids))
 
 
+# -- concurrent executions on one plugin keep their terminal records ---------
+
+
+@pytest.mark.parametrize("emit_mode", ["on-complete", "on-change"])
+def test_concurrent_executions_on_one_plugin_each_deliver_terminal_record(
+    emit_mode,
+):
+    # Two executions in flight on one plugin (as the local runner does). Each
+    # invocation end must deliver its own terminal record; a shared pending
+    # slot let one execution's RUNNING snapshot displace the other's terminal.
+    rounds = 50
+    exporter = CaptureExporter()
+    plugin = workflow_insight(
+        WorkflowInsightConfig(exporters=[exporter], emit_mode=emit_mode)
+    )
+
+    def drive(arn: str) -> None:
+        op = _step(f"{arn}-step", op_id=f"{arn}-1")
+        for _ in range(rounds):
+            plugin.on_invocation_start(_start(arn, operations={}))
+            plugin.on_operation_change(
+                OperationChangeInfo(
+                    execution_arn=arn, updated_operations=_ops(op), operations=_ops(op)
+                )
+            )
+            plugin.on_invocation_end(_end(arn, operations=_ops(op)))
+
+    threads = [threading.Thread(target=drive, args=(arn,)) for arn in (ARN, ARN_B)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(30.0)
+    assert not any(thread.is_alive() for thread in threads)
+
+    for arn in (ARN, ARN_B):
+        terminal = [
+            r
+            for r in exporter.records
+            if r["executionArn"] == arn and r["status"] == "SUCCEEDED"
+        ]
+        assert len(terminal) == rounds
+    assert plugin._state == {}
+
+
 # -- no cross-execution contamination (comment 3) ----------------------------