feat(otel): add trace-context types + resolver

Author: Ayushi Ahjolia

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

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import os
+from dataclasses import dataclass
+from enum import Enum
 from typing import TYPE_CHECKING, Callable
 
 from opentelemetry import context as otel_context, propagate
@@ -13,6 +15,45 @@
 
     from aws_durable_execution_sdk_python.plugin import InvocationStartInfo
 
+
+class Sampling(Enum):
+    """Sampling decision propagated by the durable execution backend."""
+
+    SAMPLED = "sampled"
+    NOT_SAMPLED = "not_sampled"
+    UNDECIDED = "undecided"
+
+
+@dataclass(frozen=True)
+class ExtractedContext:
+    """Trace context extracted from the durable execution backend.
+
+    Attributes:
+        trace_id: OTel 128-bit trace ID, or ``None`` when no valid trace ID was
+            present.
+        parent_span_id: OTel 64-bit parent span ID, or ``None`` when no valid
+            parent was present.
+        sampling: Explicit backend sampling decision, or ``UNDECIDED`` when
+            the backend header did not include one.
+    """
+
+    trace_id: int | None
+    parent_span_id: int | None
+    sampling: Sampling = Sampling.UNDECIDED
+
+    @property
+    def has_valid_trace_id(self) -> bool:
+        return self.trace_id is not None and 0 < self.trace_id < 2**128
+
+    @property
+    def has_valid_parent_span_id(self) -> bool:
+        return self.parent_span_id is not None and 0 < self.parent_span_id < 2**64
+
+    @property
+    def has_complete_remote_parent(self) -> bool:
+        return self.has_valid_trace_id and self.has_valid_parent_span_id
+
+
 ContextExtractor = Callable[["InvocationStartInfo"], "Context"]
 
 

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

@@ -26,9 +26,9 @@ class _IdOverride:
 def _to_otel_trace_id(execution_arn: str, start_timestamp: datetime) -> int:
     """Build a deterministic OTel-compatible execution trace ID (128 bits).
 
-    The ID is independent of ambient Lambda or X-Ray trace context so the
-    parentless Workflow span remains the only root of the durable execution
-    trace. Invocation spans inherit ambient context separately.
+    The ID is used when the backend does not provide a valid trace ID. In that
+    case a deterministic synthetic execution root anchors the durable execution
+    trace across reinvocations.
 
     Raises:
         ValueError: If the execution start timestamp is missing.
@@ -80,6 +80,22 @@ def derive_workflow_span_id(durable_execution_arn: str) -> int:
     return span_id or 1
 
 
+def derive_execution_root_span_id(durable_execution_arn: str) -> int:
+    """Derive the deterministic synthetic execution-root span ID.
+
+    The synthetic root is a non-recording parent context used when the backend
+    does not provide a complete remote parent. Its ID is stable across
+    reinvocations and uses a namespace distinct from Workflow and operation
+    span IDs.
+    """
+    if not durable_execution_arn:
+        raise ValueError("execution ARN is required to derive an execution root ID")
+    plain_value = f"execution-root:{durable_execution_arn}"
+    hashed = hashlib.blake2b(plain_value.encode()).hexdigest()[:16]
+    span_id = int(hashed, 16)
+    return span_id or 1
+
+
 class DeterministicIdGenerator(RandomIdGenerator):
     """An ID generator with invocation-scoped deterministic ID overrides.
 

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

@@ -0,0 +1,98 @@
+"""Execution trace ancestry for durable execution telemetry."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import datetime
+
+from opentelemetry.trace import SpanContext, TraceFlags, TraceState
+
+from aws_durable_execution_sdk_python_otel.context_extractors import (
+    ExtractedContext,
+    Sampling,
+)
+from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
+    _to_otel_trace_id,
+    derive_execution_root_span_id,
+)
+
+
+@dataclass(frozen=True)
+class ExecutionTraceContext:
+    """Common ancestor for Workflow and Invocation spans."""
+
+    execution_ancestor: SpanContext
+
+    @property
+    def trace_id(self) -> int:
+        return self.execution_ancestor.trace_id
+
+    @property
+    def trace_flags(self) -> TraceFlags:
+        return self.execution_ancestor.trace_flags
+
+    @classmethod
+    def resolve(
+        cls,
+        *,
+        extracted: ExtractedContext | None,
+        canonical_trace_id: int,
+        execution_arn: str,
+        root_sampled: Callable[[], bool],
+    ) -> "ExecutionTraceContext":
+        """Resolve the execution ancestor.
+
+        A complete extracted remote parent is authoritative. Otherwise a
+        deterministic synthetic root anchors all invocations of the execution on
+        the same trace.
+        """
+        sampling = extracted.sampling if extracted is not None else Sampling.UNDECIDED
+        trace_flags = _trace_flags(sampling, root_sampled)
+        if extracted is not None and extracted.has_complete_remote_parent:
+            return cls(
+                SpanContext(
+                    trace_id=canonical_trace_id,
+                    span_id=extracted.parent_span_id or 0,
+                    is_remote=True,
+                    trace_flags=trace_flags,
+                    trace_state=TraceState(),
+                )
+            )
+
+        return cls(
+            SpanContext(
+                trace_id=canonical_trace_id,
+                span_id=derive_execution_root_span_id(execution_arn),
+                is_remote=False,
+                trace_flags=trace_flags,
+                trace_state=TraceState(),
+            )
+        )
+
+
+def canonical_trace_id(
+    *,
+    extracted: ExtractedContext | None,
+    execution_arn: str,
+    execution_start_time: datetime,
+) -> int:
+    """Return the stable trace ID for this durable execution."""
+    if extracted is not None and extracted.has_valid_trace_id:
+        return extracted.trace_id or 0
+    return _to_otel_trace_id(execution_arn, execution_start_time)
+
+
+def _trace_flags(
+    sampling: Sampling,
+    root_sampled: Callable[[], bool],
+) -> TraceFlags:
+    if sampling is Sampling.SAMPLED:
+        return TraceFlags(TraceFlags.SAMPLED)
+    if sampling is Sampling.NOT_SAMPLED:
+        return TraceFlags(TraceFlags.DEFAULT)
+    return (
+        TraceFlags(TraceFlags.SAMPLED)
+        if root_sampled()
+        else TraceFlags(TraceFlags.DEFAULT)
+    )

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

@@ -12,6 +12,8 @@
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     DeterministicIdGenerator,
     _to_otel_trace_id,
+    derive_execution_root_span_id,
+    derive_workflow_span_id,
     operation_id_to_span_id,
 )
 
@@ -303,3 +305,37 @@ async def main() -> tuple[tuple[int, int], tuple[int, int]]:
 
     assert result_a == (task_a_trace_id, task_a_span_id)
     assert result_b == (task_b_trace_id, task_b_span_id)
+
+
+# ---------------------------------------------------------------------------
+# derive_execution_root_span_id
+# ---------------------------------------------------------------------------
+_ROOT_ARN = "test-arn/execution-root"
+
+
+def test_derive_execution_root_span_id_is_deterministic():
+    assert derive_execution_root_span_id(_ROOT_ARN) == derive_execution_root_span_id(
+        _ROOT_ARN
+    )
+
+
+def test_derive_execution_root_span_id_differs_by_arn():
+    assert derive_execution_root_span_id(_ROOT_ARN) != derive_execution_root_span_id(
+        _ROOT_ARN + "-other"
+    )
+
+
+def test_derive_execution_root_span_id_is_64_bit():
+    span_id = derive_execution_root_span_id(_ROOT_ARN)
+    assert 0 < span_id < 2**64
+
+
+def test_derive_execution_root_span_id_rejects_empty_arn():
+    with pytest.raises(ValueError, match="execution ARN is required"):
+        derive_execution_root_span_id("")
+
+
+def test_derive_execution_root_span_id_differs_from_workflow_span_id():
+    assert derive_execution_root_span_id(_ROOT_ARN) != derive_workflow_span_id(
+        _ROOT_ARN
+    )