fix: export visible replay completions

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

@@ -349,6 +349,31 @@ Invocation span uses that ambient span as its parent. Ambient spans on a
 different trace are ignored for durable parenting so Invocation remains on the
 execution trace.
 
+### Completion snapshots and replay
+
+The plugins consume visible terminal updates from invocation input and checkpoint
+change notifications as well as natural operation End hooks. A completed parent
+can return its cached result without revisiting an updated callback; the plugin
+exports that unvisited completion before the Invocation span closes. Records in
+the full replay snapshot alone are not new completion notifications.
+
+Pending children use a parent segment when it registers during replay, and retain
+that actual parent context if it closes before a later update arrives. Completed
+input-context metadata is used only after replay has had a chance to register its
+current segment. At invocation End, pending parents are exported before their
+children. If a parent's metadata and actual segment are unavailable, Invocation
+view falls back to Invocation; Execution view uses the derived logical parent
+context. The fallback is logged and cannot reconstruct a missing parent span.
+Operation parent links are expected to form the durable service's tree.
+
+Successful completion observations are deduplicated within one invocation.
+Invocation-view ReplayChildren and virtual contexts may still open distinct
+segments; each segment closes with its own natural outcome. Completion metadata
+and retained parent references are released at invocation End, so concurrent
+invocations and warm reuse do not share the delivery state. This is not an
+exactly-once export guarantee across failed or redelivered invocations, nor an
+acknowledgment from the telemetry backend.
+
 ### Synthetic fallback anchor
 
 Each sampled fallback invocation creates and ends `DurableExecutionRoot` at

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py

@@ -0,0 +1,65 @@
+"""Visible completion metadata owned by one instrumentation invocation."""
+
+from collections.abc import Mapping
+from graphlib import TopologicalSorter
+
+from aws_durable_execution_sdk_python.plugin import (
+    InvocationStartInfo,
+    OperationInfo,
+    OperationStatus,
+    OperationType,
+)
+
+
+_TERMINAL_STATUSES = frozenset(
+    {
+        OperationStatus.SUCCEEDED,
+        OperationStatus.FAILED,
+        OperationStatus.TIMED_OUT,
+        OperationStatus.STOPPED,
+        OperationStatus.CANCELLED,
+    }
+)
+
+
+def terminal_updates(
+    operations: Mapping[str, OperationInfo],
+) -> dict[str, OperationInfo]:
+    """Keep terminal updates, without treating replay-only state as a delivery."""
+    return {
+        operation_id: operation
+        for operation_id, operation in operations.items()
+        if operation.status in _TERMINAL_STATUSES and not operation.is_replayed
+    }
+
+
+def completed_input_parents(info: InvocationStartInfo) -> list[OperationInfo]:
+    """Visible completed contexts may parent input or later checkpoint updates."""
+    return [
+        operation
+        for operation in info.operations.values()
+        if operation.operation_type is OperationType.CONTEXT
+        and operation.status in _TERMINAL_STATUSES
+    ]
+
+
+def parent_first(
+    operations: Mapping[str, OperationInfo],
+) -> list[OperationInfo]:
+    """Materialize pending ancestors before referencing them for descendants.
+
+    Durable operations form a parent tree. Missing ancestors remain unresolved;
+    their view-specific fallback is chosen by the plugin, without inventing an
+    operation or treating its non-recording reference as an exported completion.
+    """
+    order = TopologicalSorter(
+        {
+            operation_id: (
+                (operation.parent_id,)
+                if operation.parent_id is not None and operation.parent_id in operations
+                else ()
+            )
+            for operation_id, operation in operations.items()
+        }
+    )
+    return [operations[operation_id] for operation_id in order.static_order()]

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py

@@ -42,6 +42,8 @@
     InvocationStatus,
     InvocationStartInfo,
     OperationEndInfo,
+    OperationChangeInfo,
+    OperationInfo,
     OperationStartInfo,
     OperationStatus,
     OperationType,
@@ -64,6 +66,11 @@
     TraceState,
 )
 
+from aws_durable_execution_sdk_python_otel._invocation_state import (
+    completed_input_parents,
+    parent_first,
+    terminal_updates,
+)
 from aws_durable_execution_sdk_python_otel.context_extractors import (
     ContextExtractor,
     ExtractedContext,
@@ -199,6 +206,13 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
         # paired with the thread that attached them. Every attach the plugin
         # owns is released through _detach_context so the plugin never leaves a
         # scope on the context stack.
+        # Completion delivery is invocation-owned, separate from active context
+        # segments. Registration and dequeue use the same span-registry lock.
+        self._pending_completions: dict[str, OperationInfo] = {}
+        self._pending_operation_ends: dict[str, dict[str, OperationInfo]] = {}
+        self._ending_operation_ids: set[str] = set()
+        self._input_parent_refs: dict[str, OperationInfo] = {}
+        self._completed_parent_spans: dict[str, Span] = {}
         self._context_tokens: dict[str, tuple[int, object]] = {}
         self._lock = threading.RLock()
         self._tracing_enabled = False
@@ -235,6 +249,8 @@ def _bind_sdk_tracer(self) -> bool:
     def _set_span(self, key: str, span: Span) -> None:
         with self._lock:
             self._operation_spans[key] = span
+            pending = self._pending_operation_ends.pop(key, {})
+        self._drain_registered_children(pending)
 
     def _get_span(self, key: str | None) -> Span | None:
         if key is None:
@@ -596,6 +612,10 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
         self._start_workflow_span(info)
         # Keep the invocation on the shared execution trace.
         self._start_invocation_span(info)
+        self._input_parent_refs = {
+            parent.operation_id: parent for parent in completed_input_parents(info)
+        }
+        self._remember_completions(info.updated_operations)
 
         # Make the Workflow span the active span so auto-instrumented spans
         # created during the invocation become its children. The token is
@@ -739,6 +759,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
 
     def _end_invocation_spans(self, info: InvocationEndInfo) -> None:
         """End this invocation's open spans and export the Workflow span."""
+        self._drain_pending_completions()
         # End the invocation span regardless of terminal status. Record the
         # invocation status and map it to a span status:
         #   SUCCEEDED/PENDING -> OK  (this invocation did its work, whether it
@@ -815,6 +836,13 @@ def _release_invocation_scope(self) -> None:
         self._detach_remaining_contexts()
         unbind_invocation(self)
         self._tracing_enabled = False
+        with self._lock:
+            self._pending_completions.clear()
+            self._pending_operation_ends.clear()
+            self._input_parent_refs.clear()
+            self._completed_parent_spans.clear()
+            self._ending_operation_ids.clear()
+            self._ended_operation_ids.clear()
 
     # ------------------------------------------------------------------
     # Operation lifecycle
@@ -840,24 +868,121 @@ def on_operation_start(self, info: OperationStartInfo) -> None:
         self._note_parent_start(info.parent_id, info.start_time)
 
     def on_operation_end(self, info: OperationEndInfo) -> None:
+        """Observe one natural completion, sharing delivery with snapshots."""
+        self._render_operation_end(info)
+
+    def on_operation_change(self, info: OperationChangeInfo) -> None:
+        """Retain checkpoint completions until their natural End or finalization."""
+        if self._tracing_enabled:
+            self._remember_completions(info.updated_operations)
+
+    def _remember_completions(self, updates: dict[str, OperationInfo]) -> None:
+        with self._lock:
+            for operation_id, operation in terminal_updates(updates).items():
+                if operation_id not in self._ended_operation_ids:
+                    self._pending_completions[operation_id] = operation
+                    if operation.parent_id is not None:
+                        self._pending_operation_ends.setdefault(
+                            operation.parent_id, {}
+                        )[operation_id] = operation
+
+    def _completion_parent_span(self, parent_id: str | None) -> Span | None:
+        """Prefer the active replay segment, then its retained completed context."""
+        with self._lock:
+            return (
+                self._operation_spans.get(parent_id)
+                or self._completed_parent_spans.get(parent_id)
+                if parent_id is not None
+                else self._workflow_span
+            )
+
+    def _drain_pending_completions(self) -> None:
+        """Export pending parents before children while Invocation is still open."""
+        with self._lock:
+            pending = dict(self._pending_completions)
+        for operation in parent_first(pending):
+            parent = self._completion_parent_span(operation.parent_id)
+            if parent is None:
+                parent = self._fallback_parent(operation)
+            self._render_operation_end(operation, fallback_parent=parent)
+
+    def _render_operation_end(
+        self, info: OperationInfo, *, fallback_parent: Span | None = None
+    ) -> None:
+        """Claim a completion outside customer callbacks; record success last."""
         logger.debug("Durable operation ended: %s", info)
         if not self._tracing_enabled:
             return
-        # ReplayChildren and virtual child contexts intentionally re-execute
-        # without creating a new terminal checkpoint. Their end callback is
-        # replay-only and must not re-export the deterministic logical span.
         if info.is_replayed:
             return
-        # Export the span only on the first end for this operation.
         with self._lock:
+            if info.operation_id in self._ending_operation_ids:
+                return
             if info.operation_id in self._ended_operation_ids:
                 return
-            self._ended_operation_ids.add(info.operation_id)
+            self._pending_completions[info.operation_id] = info
+            if (
+                info.parent_id is not None
+                and self._completion_parent_span(info.parent_id) is None
+                and fallback_parent is None
+            ):
+                self._pending_operation_ends.setdefault(info.parent_id, {})[
+                    info.operation_id
+                ] = info
+                return
+            self._ending_operation_ids.add(info.operation_id)
+        try:
+            self._finish_operation_end(info, fallback_parent)
+            with self._lock:
+                self._ended_operation_ids.add(info.operation_id)
+                self._pending_completions.pop(info.operation_id, None)
+                if info.parent_id is not None:
+                    siblings = self._pending_operation_ends.get(info.parent_id)
+                    if siblings is not None:
+                        siblings.pop(info.operation_id, None)
+                        if not siblings:
+                            del self._pending_operation_ends[info.parent_id]
+        finally:
+            with self._lock:
+                self._ending_operation_ids.discard(info.operation_id)
+
+    def _drain_registered_children(self, pending: dict[str, OperationInfo]) -> None:
+        """A child exporter must not interrupt its parent's context attachment."""
+        for operation in pending.values():
+            try:
+                self._render_operation_end(operation)
+            except Exception:
+                logger.exception(
+                    "Completion export failed while registering parent %s",
+                    operation.parent_id,
+                )
+
+    def _fallback_parent(self, operation: OperationInfo) -> Span:
+        parent_id = operation.parent_id
+        parent_context = (
+            self._operation_span_context(parent_id) if parent_id is not None else None
+        )
+        if parent_context is None:
+            raise ValueError(
+                "No execution trace context available for completion fallback"
+            )
+        if parent_id not in self._input_parent_refs:
+            logger.warning(
+                "Using fallback parent context for operation %s: parent %s "
+                "has no available span",
+                operation.operation_id,
+                parent_id,
+            )
+        return DurableParentSpan(parent_context)
+
+    def _finish_operation_end(
+        self, info: OperationInfo, fallback_parent: Span | None
+    ) -> None:
+        """Materialize the existing terminal span without a registry gap."""
         # An open operation is held as a non-recording placeholder; drop it and
         # create the single recording span for the operation now.
         placeholder = self._get_span(info.operation_id)
-        self._pop_span(info.operation_id)
-        parent = self._resolve_parent(info.parent_id)
+        parent = fallback_parent or self._completion_parent_span(info.parent_id)
         start_time = info.start_time
         end_time = info.end_time
         if isinstance(placeholder, DurableParentSpan):
@@ -899,9 +1024,14 @@ def on_operation_end(self, info: OperationEndInfo) -> None:
             span.set_status(StatusCode.OK)
 
         self._note_parent_end(info.parent_id, end_time)
-        popped = self._pop_span(info.operation_id)
-        if popped is not None:
-            popped.end(end_time=_to_otel_timestamp(end_time))
+        span.end(end_time=_to_otel_timestamp(end_time))
+        with self._lock:
+            if self._operation_spans.get(info.operation_id) is span:
+                del self._operation_spans[info.operation_id]
+            if info.operation_type is OperationType.CONTEXT:
+                self._completed_parent_spans[info.operation_id] = DurableParentSpan(
+                    span.get_span_context(), start_time=start_time
+                )
 
     def _start_span(
         self,
@@ -921,36 +1051,35 @@ def _start_span(
         pass ``deterministic=False`` for a fresh ID beneath the operation span.
         """
         key = span_key if span_key is not None else operation_id
-        with self._lock:
-            existing = self._operation_spans.get(key)
-            if existing is not None and existing.is_recording():
-                existing.set_attribute("durable.span.replaced_on_reentry", True)
-                existing.end()
-            links = self._build_invocation_links()
-            span_id = (
-                span_id_override
-                if span_id_override is not None
-                else (
-                    operation_id_to_span_id(self._execution_arn, operation_id)
-                    if deterministic
-                    else None
-                )
+        existing = self._get_span(key)
+        if existing is not None and existing.is_recording():
+            existing.set_attribute("durable.span.replaced_on_reentry", True)
+            existing.end()
+        links = self._build_invocation_links()
+        span_id = (
+            span_id_override
+            if span_id_override is not None
+            else (
+                operation_id_to_span_id(self._execution_arn, operation_id)
+                if deterministic
+                else None
             )
-            if parent is None:
-                parent_ctx = self._with_sampling(Context())
-            else:
-                parent_ctx = self._with_sampling(
-                    trace.set_span_in_context(parent, Context())
-                )
-            with self._id_generator.use_ids(trace_id=None, span_id=span_id):
-                span = self._tracer.start_span(
-                    name=name,
-                    attributes=self._operation_attributes(info),
-                    start_time=_to_otel_timestamp(start_time),
-                    context=parent_ctx,
-                    links=links,
-                )
-            self._operation_spans[key] = span
+        )
+        if parent is None:
+            parent_ctx = self._with_sampling(Context())
+        else:
+            parent_ctx = self._with_sampling(
+                trace.set_span_in_context(parent, Context())
+            )
+        with self._id_generator.use_ids(trace_id=None, span_id=span_id):
+            span = self._tracer.start_span(
+                name=name,
+                attributes=self._operation_attributes(info),
+                start_time=_to_otel_timestamp(start_time),
+                context=parent_ctx,
+                links=links,
+            )
+        self._set_span(key, span)
         return span
 
     # ------------------------------------------------------------------