fix(otel): preserve replay completion telemetry

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

@@ -405,13 +405,108 @@ each durable span in the same invocation.
 
 `durable.invocation.status` uses `RETRYING` when the core plugin hook reports
 `InvocationStatus.RETRY`. The core enum remains unchanged. Operation spans use
-OTel `OK` only for `SUCCEEDED`, `ERROR` when error details are delivered, and
-`UNSET` for other outcomes without error details, including `FAILED`,
-`CANCELLED`, `TIMED_OUT`, and `STOPPED`. The original durable operation status
-remains in `durable.operation.status`.
+OTel `ERROR` for explicit `FAILED`, `TIMED_OUT`, or `STOPPED` outcomes even when
+error details are absent. Supplied error details continue to set `ERROR` and
+record an exception event; a failure status alone does not invent that event.
+Without error details, `SUCCEEDED` is `OK` and other statuses, including
+`CANCELLED`, remain `UNSET`. The original durable operation status remains in
+`durable.operation.status`. A handler that catches an operation failure and
+returns successfully still has a successful invocation/workflow outcome; the
+failed operation retains its own error status on initial completion and replay.
+
+Exception events use the same normalized end timestamp as their operation or
+attempt span. This keeps events inside the exported span when a completion has
+no durable start time or arrives later than its backend end time; the existing
+view-specific span-boundary rules remain intact.
+
+When a terminal completion arrives before replay creates its parent span, the
+plugin retains it until that span is registered. Parent registration exports the
+pending child inside the actual parent segment. Completions whose parents remain
+unavailable are exported at invocation End: Invocation view uses the invocation
+span, and Execution view derives a logical parent context from the operation ID.
+Each fallback is logged. A derived context is not a reconstructed recording span
+or proof of an unavailable virtual context's original span identity.
+
+### Invocation flushing
+
+After span and ContextVar cleanup, the plugins wait up to the OpenTelemetry
+provider's existing 30-second default budget for `force_flush()`. Only provider
+I/O runs on a background daemon worker; invocation Start/End hooks and token
+cleanup stay on their original worker and Context. The I/O receives a snapshot
+of that post-cleanup Context.
+
+There is at most one outstanding flush per provider and no pending I/O queue.
+Distinct providers can flush independently. Concurrent requests for the same
+provider wait for admission within the same absolute deadline used for completion,
+then perform a fresh flush so spans added during the earlier attempt are covered.
+A timeout does not cancel the exporter or free that provider's slot. A request
+whose budget expires before admission returns with a warning without starting
+more work.
+A failed or skipped flush is not an export
+acknowledgment, and late exports may continue. Applications still own exporter
+transport timeouts and provider shutdown. This bound does not cover synchronous
+span-processor `on_end` work that runs before flushing.
+
+Ordinary provider exceptions that finish within the budget retain the existing
+plugin exception isolation; fatal exceptions are propagated on the lifecycle
+owner. Exceptions after the caller has stopped waiting are logged and cannot be
+raised retroactively into the completed invocation.
 
 ### Invocation context isolation
 
+Factories create a separate plugin instance for each invocation. Each instance
+owns its trace IDs, span registry, deferred completions, timestamps and scope
+tokens, so one invocation's cleanup cannot clear another's trace. Factories share
+provider configuration across instances. SDK user-code workers capture the
+invocation Context; Start/End hooks and token reset retain their owner and order.
+
+Each OTel plugin retains visible terminal operations from
+`InvocationStartInfo.updated_operations` and checkpoint `on_operation_change`
+snapshots. Natural `on_operation_end` hooks export those operations and remove
+them from the pending dictionary. Pending children are also indexed by parent,
+so replay registration can export them under the actual parent before it closes.
+The plugin exports any unvisited completions before ending the invocation span.
+An invocation-local set suppresses repeated input, replay and checkpoint exports;
+core does not synthesize additional completion hooks for other plugins.
+
+A completion enters that set only after its span ends successfully. An in-progress
+claim prevents concurrent attempts, and a failed attempt leaves the completion
+pending for a later notification or invocation-End retry. If a processor delivers
+a span and then raises, retrying can repeat that delivery; the plugin has no
+exporter acknowledgment that would distinguish partial delivery from failure.
+
+During invocation End, an ordinary completion-export failure receives at most one
+retry. Failures do not prevent other completions or the Invocation/Workflow spans
+from finishing. Errors are reported after scope cleanup and flushing, including
+recovered transient failures. Control exceptions are not retried and propagate
+after cleanup; they are not converted into ordinary export failures.
+
+Invocation view can open another context segment within the same invocation
+when a timed branch resumes and re-enters a completed child, or when a virtual
+context runs again. Each new segment receives its own terminal outcome. Repeated
+durable checkpoint observations still do not close or re-export that segment.
+
+The service omits descendants of completed contexts unless their ancestors retain
+children through `ReplayChildren`. An `UpdatedOperationIds` entry alone does not
+provide missing operation metadata, and the plugin does not synthesize an
+operation span from that missing record.
+
+Known completed context metadata is retained from the full visible input snapshot
+for descendants delivered by later checkpoint responses. Both views retain actual
+completed parent contexts until invocation cleanup. Invocation view prefers its
+current replay segment, including a large-result ReplayChildren parent; only at
+End can a completion without such a segment use a known completed input parent's
+identity. Remaining unknown parents use the logged fallback described above.
+Queue draining does not re-enter public hooks.
+
+Completion metadata can be redelivered after a failed invocation or checkpoint
+transport attempt that has no successful response acknowledgment. There is no
+exporter acknowledgment in durable state, so this does not promise exactly-once
+telemetry across failed/redelivered invocations. Successful replay preserves
+checkpointed work; this limit is not a claim that already-checkpointed step
+bodies run again. The plugin does not use a global exported-operation cache that
+could discard a first delivery after a failed attempt.
+
 With core 3.x, invocation hooks and the handler run on one worker in an
 invocation-local Context initialized from the host's bindings. Start bindings
 remain visible to later hooks, the handler and resource cleanup even when Start

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

@@ -0,0 +1,128 @@
+"""Bound caller waiting for provider I/O without moving plugin lifecycle hooks."""
+
+from __future__ import annotations
+
+import contextvars
+import logging
+import threading
+import time
+from dataclasses import dataclass, field
+
+from opentelemetry.trace import TracerProvider
+
+
+logger = logging.getLogger(__name__)
+
+# Match the OTel provider's existing default budget. Some processors ignore
+# that cooperative budget, so the SDK bounds its own wait independently.
+_FLUSH_TIMEOUT_SECONDS = 30.0
+_lock = threading.Condition(threading.Lock())
+
+
+@dataclass
+class _Attempt:
+    done: threading.Event = field(default_factory=threading.Event)
+    result: object = None
+    error: BaseException | None = None
+    timed_out: bool = False
+
+
+# Identity keys support custom providers with value equality or no hash. The
+# worker retains its provider until removal, so an active key cannot be reused.
+_active: dict[int, _Attempt] = {}
+
+
+def _run(provider: TracerProvider, attempt: _Attempt) -> None:
+    provider_id = id(provider)
+    try:
+        try:
+            # force_flush is an optional extension to the API provider. The
+            # plugin checks for it before invoking this helper.
+            attempt.result = getattr(provider, "force_flush")()
+        except BaseException as error:
+            # Fatal exceptions completed within the wait are re-raised on the
+            # lifecycle owner. A late error cannot be thrown into a returned call.
+            attempt.error = error
+        with _lock:
+            late = attempt.timed_out
+            if not late and _active.get(provider_id) is attempt:
+                del _active[provider_id]
+                _lock.notify_all()
+            attempt.done.set()
+        if late and attempt.error is not None:
+            logger.error(
+                "Provider flush failed after the caller stopped waiting",
+                exc_info=(
+                    type(attempt.error),
+                    attempt.error,
+                    attempt.error.__traceback__,
+                ),
+            )
+    finally:
+        # Retain the slot through actual completion, even after timeout. There
+        # is no queue and no new worker per invocation for a stuck exporter.
+        with _lock:
+            if _active.get(provider_id) is attempt:
+                del _active[provider_id]
+                _lock.notify_all()
+
+
+def flush_provider(provider: TracerProvider) -> bool:
+    """Wait for provider I/O, with one outstanding attempt per provider.
+
+    Start/End and ContextVar-token cleanup stay on their existing owner. Only
+    the supplied provider method runs elsewhere, using the post-cleanup Context
+    snapshot. Admission and completion share one absolute deadline. Expired
+    requests return False without queueing more work; active I/O is never cancelled.
+    Each admitted caller starts a new flush after prior I/O completes, since spans
+    queued during a prior attempt may not have been included in that attempt.
+    This does not bound synchronous span processors' earlier ``on_end`` calls,
+    application-owned provider shutdown, or exporter transport I/O itself.
+    """
+    deadline = time.monotonic() + _FLUSH_TIMEOUT_SECONDS
+    provider_id = id(provider)
+    with _lock:
+        while provider_id in _active:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                break
+            _lock.wait(remaining)
+        if provider_id not in _active and time.monotonic() < deadline:
+            attempt = _Attempt()
+            _active[provider_id] = attempt
+        else:
+            attempt = None
+    if attempt is None:
+        logger.warning("Provider flush admission exceeded the caller wait budget")
+        return False
+
+    owner_context = contextvars.copy_context()
+    try:
+        threading.Thread(
+            target=owner_context.run,
+            args=(_run, provider, attempt),
+            name="durable-otel-flush",
+            daemon=True,
+        ).start()
+    except BaseException:
+        with _lock:
+            if _active.get(provider_id) is attempt:
+                del _active[provider_id]
+                _lock.notify_all()
+        raise
+
+    completed = attempt.done.wait(max(0.0, deadline - time.monotonic()))
+    if not completed:
+        with _lock:
+            completed = attempt.done.is_set()
+            if not completed:
+                attempt.timed_out = True
+        if not completed:
+            logger.warning("Provider flush exceeded the caller wait budget")
+            return False
+    if attempt.error is not None:
+        raise attempt.error
+    if attempt.result is False:
+        logger.warning("Provider reported an incomplete flush")
+        return False
+    return True

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

@@ -0,0 +1,92 @@
+"""Invocation-owned completion metadata and bounded finalization."""
+
+from collections.abc import Callable, Mapping
+
+from aws_durable_execution_sdk_python.plugin import (
+    InvocationStartInfo,
+    OperationInfo,
+    OperationStatus,
+    OperationType,
+)
+
+
+class InvocationFinalization:
+    """Defer failures until invocation cleanup, with one ordinary export retry.
+
+    Operation claims and failures use the plugin's registry lock. Finalization
+    actions and error reporting run on the invocation owner. Control exceptions
+    are held only to finish cleanup; they are never retried or swallowed.
+    """
+
+    def __init__(self) -> None:
+        self._attempts: dict[str, int] = {}
+        self._stopped_operations: set[str] = set()
+        self._retryable_operations: set[str] = set()
+        self._errors: list[BaseException] = []
+
+    def claim_operation(self, operation_id: str) -> bool:
+        attempts = self._attempts.get(operation_id, 0)
+        if attempts >= 2 or operation_id in self._stopped_operations:
+            return False
+        self._attempts[operation_id] = attempts + 1
+        self._retryable_operations.discard(operation_id)
+        return True
+
+    def operation_failed(self, operation_id: str, error: BaseException) -> None:
+        self._errors.append(error)
+        if not isinstance(error, Exception):
+            self._stopped_operations.add(operation_id)
+        elif self._attempts[operation_id] < 2:
+            self._retryable_operations.add(operation_id)
+
+    def can_retry(self, operation_id: str) -> bool:
+        return operation_id in self._retryable_operations
+
+    def run(self, action: Callable[[], None]) -> None:
+        try:
+            action()
+        except BaseException as error:
+            self._errors.append(error)
+
+    def raise_errors(self) -> None:
+        if len(self._errors) == 1:
+            raise self._errors[0]
+        if self._errors:
+            raise BaseExceptionGroup(
+                "OpenTelemetry invocation finalization failed", self._errors
+            )
+
+
+def terminal_updates(
+    operations: Mapping[str, OperationInfo],
+) -> dict[str, OperationInfo]:
+    """Retain visible terminal updates until their operation hook or invocation End."""
+    return {
+        operation_id: operation
+        for operation_id, operation in operations.items()
+        if operation.status
+        in (
+            OperationStatus.SUCCEEDED,
+            OperationStatus.FAILED,
+            OperationStatus.TIMED_OUT,
+            OperationStatus.STOPPED,
+            OperationStatus.CANCELLED,
+        )
+    }
+
+
+def completed_input_parents(info: InvocationStartInfo) -> list[OperationInfo]:
+    """Known terminal contexts that 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 (
+            OperationStatus.SUCCEEDED,
+            OperationStatus.FAILED,
+            OperationStatus.TIMED_OUT,
+            OperationStatus.STOPPED,
+            OperationStatus.CANCELLED,
+        )
+    ]