fix(otel): stop reusing operation span IDs on resume

The execution plugin now ends operation spans left open at invocation
end, so an operation that suspends and resumes across invocations would
export its deterministic operation span ID more than once, corrupting
the trace. Mirror the invocation plugin: a replayed or cross-invocation
operation span uses a fresh span ID and links back to the deterministic
operation context, so each per-invocation segment correlates without
sharing a span ID.

Also resolve the Workflow placeholder and operation link trace state
from the sampling result (falling back to the ancestor state) in both
plugins, so a same-trace ambient tracestate is preserved on non-terminal
invocations instead of being dropped to the empty ancestor state.

sim: #642

Author: Ayushi Ahjolia

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

@@ -56,6 +56,7 @@
     SpanKind,
     StatusCode,
     Tracer,
+    TraceState,
 )
 
 from aws_durable_execution_sdk_python_otel.context_extractors import (
@@ -293,6 +294,41 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None:
                 return existing
         return self._workflow_span
 
+    def _resolved_trace_state(self) -> TraceState:
+        """Return the trace state shared by every durable span this invocation.
+
+        The sampling result carries the ``tracestate`` resolved for the
+        execution: when the invocation runs under a same-trace ambient span the
+        result preserves that vendor state, whereas the execution ancestor's
+        state is empty. Prefer the resolved state so placeholders match the
+        recording spans; fall back to the ancestor state when no intent exists.
+        """
+        intent = self._sampling_intent
+        if intent is not None and intent.result.trace_state is not None:
+            return intent.result.trace_state
+        if self._execution_trace_context is not None:
+            return self._execution_trace_context.execution_ancestor.trace_state
+        return TraceState()
+
+    def _operation_link_context(self, operation_id: str) -> SpanContext | None:
+        """Return the deterministic logical operation context for links.
+
+        A continuation span (a replayed operation, or one stitched across
+        invocations) uses a fresh span ID and links back to this deterministic
+        context so every per-invocation segment of one logical operation
+        correlates without sharing a span ID.
+        """
+        execution_trace_context = self._execution_trace_context
+        if execution_trace_context is None:
+            return None
+        return SpanContext(
+            trace_id=execution_trace_context.trace_id,
+            span_id=operation_id_to_span_id(self._execution_arn, operation_id),
+            is_remote=False,
+            trace_flags=execution_trace_context.trace_flags,
+            trace_state=self._resolved_trace_state(),
+        )
+
     def _invocation_parent_context(self) -> Context:
         """Return same-trace ambient context, else execution ancestor context."""
         execution_trace_context = self._execution_trace_context
@@ -415,7 +451,7 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None:
             span_id=derive_workflow_span_id(self._execution_arn),
             is_remote=False,
             trace_flags=self._execution_trace_context.trace_flags,
-            trace_state=self._execution_trace_context.execution_ancestor.trace_state,
+            trace_state=self._resolved_trace_state(),
         )
         self._workflow_span = NonRecordingSpan(workflow_span_context)
 
@@ -568,6 +604,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None:
             info=info,
             parent=parent,
             start_time=info.start_time,
+            existed=info.is_replayed,
         )
 
     def on_operation_end(self, info: OperationEndInfo) -> None:
@@ -577,14 +614,16 @@ def on_operation_end(self, info: OperationEndInfo) -> None:
         span = self._get_span(info.operation_id)
         if span is None:
             # Cross-invocation stitching: operation started in a prior
-            # invocation. Create + immediately end a linked span.
+            # invocation. Create + immediately end a continuation span with a
+            # fresh ID that links back to the deterministic operation context.
             parent = self._resolve_parent(info.parent_id)
             span = self._start_span(
                 operation_id=info.operation_id,
                 name=info.name or info.operation_id,
                 info=info,
                 parent=parent,
                 start_time=info.start_time,
+                existed=True,
             )
         else:
             span.set_attributes(self._operation_attributes(info))
@@ -614,16 +653,29 @@ def _start_span(
         start_time: datetime.datetime | None,
         span_key: str | None = None,
         deterministic: bool = True,
+        existed: bool = False,
     ) -> Span:
-        """Start a span for an operation/attempt and register it."""
+        """Start a span for an operation/attempt and register it.
+
+        The first span observed for a logical operation uses the deterministic
+        operation span ID. A continuation -- a replayed operation, or one
+        stitched across invocations in ``on_operation_end`` -- instead uses a
+        fresh span ID and links back to the deterministic operation context, so
+        the segments correlate without two exported spans sharing one ID.
+        """
         key = span_key if span_key is not None else operation_id
         with self._lock:
             links = self._build_invocation_links()
+            use_deterministic_id = deterministic and not existed
             span_id = (
                 operation_id_to_span_id(self._execution_arn, operation_id)
-                if deterministic
+                if use_deterministic_id
                 else None
             )
+            if deterministic and existed:
+                operation_context = self._operation_link_context(operation_id)
+                if operation_context is not None and operation_context.is_valid:
+                    links = [*links, Link(context=operation_context)]
 
             if parent is None:
                 parent_ctx = self._with_sampling(Context())

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

@@ -31,6 +31,7 @@
     SpanKind,
     StatusCode,
     Tracer,
+    TraceState,
 )
 
 from aws_durable_execution_sdk_python_otel.context_extractors import (
@@ -359,6 +360,23 @@ def _next_ordered_timestamp(
             self._span_time_floor_ns = candidate
         return candidate
 
+    def _resolved_trace_state(self) -> TraceState:
+        """Return the trace state shared by every durable span this invocation.
+
+        The sampling result carries the ``tracestate`` resolved for the
+        execution: when the invocation runs under a same-trace ambient span the
+        result preserves that vendor state, whereas the execution ancestor's
+        state is empty. Prefer the resolved state so placeholders and links
+        match the recording spans; fall back to the ancestor state when no
+        intent exists.
+        """
+        intent = self._sampling_intent
+        if intent is not None and intent.result.trace_state is not None:
+            return intent.result.trace_state
+        if self._execution_trace_context is not None:
+            return self._execution_trace_context.execution_ancestor.trace_state
+        return TraceState()
+
     def _operation_link_context(self, operation_id: str) -> SpanContext | None:
         """Return the deterministic logical operation context for links."""
         execution_trace_context = self._execution_trace_context
@@ -369,7 +387,7 @@ def _operation_link_context(self, operation_id: str) -> SpanContext | None:
             span_id=operation_id_to_span_id(self._execution_arn, operation_id),
             is_remote=False,
             trace_flags=execution_trace_context.trace_flags,
-            trace_state=execution_trace_context.execution_ancestor.trace_state,
+            trace_state=self._resolved_trace_state(),
         )
 
     def _start_span(
@@ -573,7 +591,7 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None:
             span_id=derive_workflow_span_id(self._execution_arn),
             is_remote=False,
             trace_flags=self._execution_trace_context.trace_flags,
-            trace_state=self._execution_trace_context.execution_ancestor.trace_state,
+            trace_state=self._resolved_trace_state(),
         )
         self._workflow_span = NonRecordingSpan(workflow_span_context)
 

packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py

@@ -26,11 +26,19 @@
     UserFunctionStartInfo,
 )
 from opentelemetry import baggage, trace
+from opentelemetry.context import Context
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import (
+    NonRecordingSpan,
+    SpanContext,
+    TraceFlags,
+    TraceState,
+)
 
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
+    _to_otel_trace_id,
     derive_execution_root_span_id,
     derive_workflow_span_id,
     operation_id_to_span_id,
@@ -343,14 +351,20 @@ def test_operation_parented_under_workflow_and_linked_to_invocation():
     assert invocation.context.span_id in linked_span_ids
 
 
-def test_cross_invocation_operation_end_uses_deterministic_span_id():
+def test_cross_invocation_operation_end_links_previous_logical_operation():
+    """A continuation uses a fresh ID and links the deterministic operation."""
     plugin, exporter = _create_plugin()
     plugin.on_invocation_start(_invocation_start_info())
+    operation_id = "step-earlier"
+    random_span_id = int("1234567890abcdef", 16)
+    plugin._id_generator._fallback_id_generator.generate_span_id = lambda: (
+        random_span_id
+    )
 
     # Backend-updated completion for an operation started in a prior invocation.
     plugin.on_operation_end(
         OperationEndInfo(
-            operation_id="step-earlier",
+            operation_id=operation_id,
             operation_type=OperationType.STEP,
             sub_type=None,
             name="earlier-step",
@@ -365,12 +379,16 @@ def test_cross_invocation_operation_end_uses_deterministic_span_id():
     plugin.on_invocation_end(_invocation_end_info())
 
     matching = [s for s in exporter.get_finished_spans() if s.name == "earlier-step"]
-    # Exported exactly once, using the deterministic logical-operation span ID
-    # (no separate continuation span).
+    # Exported exactly once, with a fresh span ID (not the deterministic one) so
+    # a later terminal completion cannot collide with this segment.
     assert len(matching) == 1
-    assert matching[0].context.span_id == operation_id_to_span_id(
-        EXECUTION_ARN, "step-earlier"
+    assert matching[0].context.span_id == random_span_id
+    assert matching[0].context.span_id != operation_id_to_span_id(
+        EXECUTION_ARN, operation_id
     )
+    # Links back to the deterministic logical operation context for correlation.
+    linked_span_ids = {link.context.span_id for link in matching[0].links}
+    assert operation_id_to_span_id(EXECUTION_ARN, operation_id) in linked_span_ids
 
 
 @pytest.mark.parametrize(
@@ -621,6 +639,144 @@ def test_open_operation_span_ended_at_invocation_end():
     assert "wait-for-signal" in exported
 
 
+def test_suspend_then_resume_operation_exports_unique_span_ids():
+    """An operation that suspends then resumes never reuses a span ID.
+
+    Invocation N starts the operation and suspends: its deterministic-ID span
+    is ended at invocation end. Invocation N+1 replays the still-open operation
+    (on_operation_replay -> on_operation_start with is_replayed=True) and then
+    completes it. Each exported segment must carry a distinct span ID so no two
+    exported spans collide on one span_id/trace_id.
+    """
+    plugin, exporter = _create_plugin()
+    operation_id = "wait-across-invocations"
+    replay_span_id = int("1111111111111111", 16)
+    resume_span_id = int("2222222222222222", 16)
+
+    # Invocation N: operation starts and suspends (non-terminal invocation end).
+    plugin.on_invocation_start(_invocation_start_info())
+    plugin.on_operation_start(
+        OperationStartInfo(
+            operation_id=operation_id,
+            operation_type=OperationType.WAIT,
+            sub_type=OperationSubType.WAIT,
+            name="long-wait",
+            parent_id=None,
+            start_time=START_TIME,
+            is_replayed=False,
+            status=OperationStatus.STARTED,
+        )
+    )
+    plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING))
+
+    # Invocation N+1: the still-open operation is replayed, then completes.
+    plugin.on_invocation_start(_invocation_start_info())
+    plugin._id_generator._fallback_id_generator.generate_span_id = lambda: (
+        replay_span_id
+    )
+    plugin.on_operation_start(
+        OperationStartInfo(
+            operation_id=operation_id,
+            operation_type=OperationType.WAIT,
+            sub_type=OperationSubType.WAIT,
+            name="long-wait",
+            parent_id=None,
+            start_time=START_TIME,
+            is_replayed=True,
+            status=OperationStatus.STARTED,
+        )
+    )
+    plugin._id_generator._fallback_id_generator.generate_span_id = lambda: (
+        resume_span_id
+    )
+    plugin.on_operation_end(
+        OperationEndInfo(
+            operation_id=operation_id,
+            operation_type=OperationType.WAIT,
+            sub_type=OperationSubType.WAIT,
+            name="long-wait",
+            parent_id=None,
+            start_time=START_TIME,
+            is_replayed=False,
+            status=OperationStatus.SUCCEEDED,
+            end_time=END_TIME,
+            error=None,
+        )
+    )
+    plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.SUCCEEDED))
+
+    operation_spans = [
+        s for s in exporter.get_finished_spans() if s.name == "long-wait"
+    ]
+    span_ids = [s.context.span_id for s in operation_spans]
+    # Every exported segment has a unique span ID (no reuse across invocations).
+    assert len(span_ids) == len(set(span_ids))
+    # The first segment used the deterministic ID; every continuation used a
+    # fresh ID and links back to the deterministic logical operation context.
+    deterministic_id = operation_id_to_span_id(EXECUTION_ARN, operation_id)
+    assert deterministic_id in span_ids
+    for span in operation_spans:
+        if span.context.span_id != deterministic_id:
+            linked = {link.context.span_id for link in span.links}
+            assert deterministic_id in linked
+
+
+def test_pre_terminal_placeholder_preserves_same_trace_tracestate():
+    """The Workflow placeholder and operation links carry ambient tracestate.
+
+    A same-trace ambient span carries vendor ``tracestate``. On a non-terminal
+    invocation the recording Workflow span is never exported, so the placeholder
+    is the only Workflow context that parents operation spans; it (and the
+    deterministic operation link a continuation emits) must carry that resolved
+    tracestate rather than the empty ancestor state.
+    """
+    plugin, exporter = _create_plugin()
+    canonical = _to_otel_trace_id(EXECUTION_ARN, START_TIME)
+    trace_state = TraceState([("vendor", "opaque")])
+    ambient_context = SpanContext(
+        trace_id=canonical,
+        span_id=int("1234567890abcdef", 16),
+        is_remote=False,
+        trace_flags=TraceFlags(TraceFlags.SAMPLED),
+        trace_state=trace_state,
+    )
+    ambient = NonRecordingSpan(ambient_context)
+    token = otel_context.attach(trace.set_span_in_context(ambient, Context()))
+    try:
+        plugin.on_invocation_start(_invocation_start_info())
+        # Placeholder Workflow span carries the resolved tracestate.
+        assert plugin._workflow_span is not None
+        assert plugin._workflow_span.get_span_context().trace_state == trace_state
+        # A cross-invocation completion links the deterministic operation
+        # context, which must also carry the resolved tracestate.
+        plugin.on_operation_end(
+            OperationEndInfo(
+                operation_id="wait-existing",
+                operation_type=OperationType.WAIT,
+                sub_type=OperationSubType.WAIT,
+                name="existing-wait",
+                parent_id=None,
+                start_time=START_TIME,
+                is_replayed=False,
+                status=OperationStatus.SUCCEEDED,
+                end_time=END_TIME,
+                error=None,
+            )
+        )
+        plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING))
+    finally:
+        otel_context.detach(token)
+
+    span = next(s for s in exporter.get_finished_spans() if s.name == "existing-wait")
+    operation_link = next(
+        link
+        for link in span.links
+        if link.context.span_id
+        == operation_id_to_span_id(EXECUTION_ARN, "wait-existing")
+    )
+    assert operation_link.context.trace_state == trace_state
+
+
 @pytest.mark.parametrize(
     ("status", "expected_code"),
     [