feat: add otel logger injection

Author: Alex Wang

packages/aws-durable-execution-sdk-python-examples/examples-catalog.json

@@ -635,6 +635,21 @@
         "ExecutionTimeout": 300
       },
       "path": "./src/plugin/execution_with_otel.py"
+    },
+    {
+      "name": "Otel Logger Example",
+      "description": "Demonstrates OTel-enriched logging correlated to durable spans",
+      "handler": "otel_logger_example.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "loggingConfig": {
+        "ApplicationLogLevel": "INFO",
+        "LogFormat": "JSON"
+      },
+      "path": "./src/otel/otel_logger_example.py"
     }
   ]
 }

packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py

@@ -0,0 +1,67 @@
+"""Demonstrates OTel-enriched logging in a durable execution.
+
+The DurableExecutionOtelPlugin wraps the execution logger (enrich_logger=True
+by default) so every log line emitted through context.logger / step_context.logger
+is automatically enriched with the active OpenTelemetry trace context
+(otel.trace_id, otel.span_id, otel.trace_sampled). This lets logs correlate to
+the spans the plugin emits without any user code changes.
+
+Logs emitted:
+- at the top level correlate to the invocation span
+- inside a step correlate to that step's span
+- inside a child context correlate to the child-context span
+"""
+
+from typing import Any
+
+from aws_durable_execution_sdk_python_otel import DurableExecutionOtelPlugin
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+
+from aws_durable_execution_sdk_python import StepContext
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    durable_step,
+    durable_with_child_context,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+
+
+tracer_provider = TracerProvider()
+trace.set_tracer_provider(tracer_provider)
+
+# enrich_logger defaults to True, so the execution logger is wrapped with OTel
+# trace context injection (otel.trace_id, otel.span_id, otel.trace_sampled).
+otel = DurableExecutionOtelPlugin(tracer_provider)
+
+
+@durable_step
+def greet(step_context: StepContext, name: str) -> str:
+    # Logged inside a step: enriched with this step's span_id.
+    # Note: avoid reserved LogRecord keys (e.g. "name") in extra.
+    step_context.logger.info("Greeting inside step", extra={"greeting_name": name})
+    return f"hello {name}"
+
+
+@durable_with_child_context
+def greet_in_child(child_context: DurableContext, name: str) -> str:
+    # Logged inside a child context: enriched with the child-context span_id.
+    child_context.logger.info("Entering child context")
+    result: str = child_context.step(greet(name), name="child-greet")
+    child_context.logger.info("Leaving child context", extra={"result": result})
+    return result
+
+
+@durable_execution(plugins=[otel])
+def handler(_event: Any, context: DurableContext) -> str:
+    # Logged at the top level: enriched with the invocation span_id.
+    context.logger.info("Workflow started")
+
+    top: str = context.step(greet("world"), name="top-greet")
+    nested: str = context.run_in_child_context(
+        greet_in_child("nested"), name="child-context"
+    )
+
+    context.logger.info("Workflow completed", extra={"top": top, "nested": nested})
+    return f"{top} | {nested}"

packages/aws-durable-execution-sdk-python-examples/template.yaml

@@ -1031,6 +1031,24 @@
           "ExecutionTimeout": 300
         }
       }
+    },
+    "OtelLoggerExample": {
+      "Type": "AWS::Serverless::Function",
+      "Properties": {
+        "CodeUri": "build/",
+        "Handler": "otel_logger_example.handler",
+        "Description": "Demonstrates OTel-enriched logging correlated to durable spans",
+        "Role": {
+          "Fn::GetAtt": [
+            "DurableFunctionRole",
+            "Arn"
+          ]
+        },
+        "DurableConfig": {
+          "RetentionPeriodInDays": 7,
+          "ExecutionTimeout": 300
+        }
+      }
     }
   }
 }

packages/aws-durable-execution-sdk-python-examples/test/otel/test_otel_logger_example.py

@@ -0,0 +1,32 @@
+"""Tests for the OTel-enriched logger example."""
+
+import pytest
+
+from aws_durable_execution_sdk_python.execution import InvocationStatus
+from aws_durable_execution_sdk_python.lambda_service import OperationType
+from src.otel import otel_logger_example
+from test.conftest import deserialize_operation_payload
+
+
+@pytest.mark.example
+@pytest.mark.durable_execution(
+    handler=otel_logger_example.handler,
+    lambda_function_name="Otel Logger Example",
+)
+def test_otel_logger_example(durable_runner):
+    """Verify the OTel logger example runs and produces the expected result."""
+    with durable_runner:
+        result = durable_runner.run(input="{}", timeout=10)
+
+    assert result.status is InvocationStatus.SUCCEEDED
+    assert deserialize_operation_payload(result.result) == "hello world | hello nested"
+
+    # The top-level step is named "top-greet".
+    top_step = result.get_step("top-greet")
+    assert deserialize_operation_payload(top_step.result) == "hello world"
+
+    # The child context wraps a nested step, so a CONTEXT operation exists.
+    context_ops = [
+        op for op in result.operations if op.operation_type is OperationType.CONTEXT
+    ]
+    assert len(context_ops) >= 1

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

@@ -9,6 +9,7 @@
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     DeterministicIdGenerator,
 )
+from aws_durable_execution_sdk_python_otel.logger import OtelEnrichedLogger
 from aws_durable_execution_sdk_python_otel.plugin import (
     DurableExecutionOtelPlugin,
 )
@@ -19,6 +20,7 @@
     "ContextExtractor",
     "DeterministicIdGenerator",
     "DurableExecutionOtelPlugin",
+    "OtelEnrichedLogger",
     "w3c_client_context_extractor",
     "xray_context_extractor",
 ]

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

@@ -0,0 +1,86 @@
+"""OTel-enriched logger for durable executions.
+
+Provides a LoggerInterface wrapper that injects OpenTelemetry trace context
+(trace_id, span_id, trace_sampled) into every log message's extra dict. This
+enables log-trace correlation in observability backends without changing user code.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING
+
+from opentelemetry.trace import TraceFlags
+
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.types import LoggerInterface
+
+    from aws_durable_execution_sdk_python_otel.plugin import DurableExecutionOtelPlugin
+
+
+class OtelEnrichedLogger:
+    """LoggerInterface wrapper that injects OTel trace context into log extra fields.
+
+    The span context is resolved by the plugin via get_current_span_context(),
+    which returns the active operation span inside steps and the invocation span
+    for top-level handler code.
+
+    Injected fields:
+        - otel.trace_id: 32-char hex trace identifier
+        - otel.span_id: 16-char hex span identifier
+        - otel.trace_sampled: boolean indicating if the trace is sampled
+
+    Args:
+        inner: The underlying logger to delegate to after enrichment.
+        plugin: The OTel plugin instance that resolves the current span context.
+    """
+
+    def __init__(
+        self, inner: LoggerInterface, plugin: DurableExecutionOtelPlugin
+    ) -> None:
+        self._inner = inner
+        self._plugin = plugin
+
+    def debug(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._inner.debug(msg, *args, extra=self._enrich(extra))
+
+    def info(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._inner.info(msg, *args, extra=self._enrich(extra))
+
+    def warning(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._inner.warning(msg, *args, extra=self._enrich(extra))
+
+    def error(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._inner.error(msg, *args, extra=self._enrich(extra))
+
+    def exception(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None:
+        self._inner.exception(msg, *args, extra=self._enrich(extra))
+
+    def _enrich(self, extra: Mapping[str, object] | None) -> dict[str, object]:
+        """Inject OTel trace context into the extra dict.
+
+        trace_id, span_id, and trace_sampled come from the span context resolved
+        by the plugin, so the values always match the exported spans.
+        """
+        enriched: dict[str, object] = dict(extra) if extra else {}
+
+        span_context = self._plugin.get_current_span_context()
+        if span_context and span_context.is_valid:
+            enriched["otel.trace_id"] = format(span_context.trace_id, "032x")
+            enriched["otel.span_id"] = format(span_context.span_id, "016x")
+            enriched["otel.trace_sampled"] = bool(
+                span_context.trace_flags & TraceFlags.SAMPLED
+            )
+
+        return enriched

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

@@ -20,7 +20,6 @@
 )
 from opentelemetry import context, trace
 from opentelemetry.context import Context
-from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider
 from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
 from opentelemetry.trace import (
@@ -40,10 +39,11 @@
     DeterministicIdGenerator,
     operation_id_to_span_id,
 )
+from aws_durable_execution_sdk_python_otel.logger import OtelEnrichedLogger
 
 
 if TYPE_CHECKING:
-    pass
+    from aws_durable_execution_sdk_python.types import LoggerInterface
 
 
 logger = logging.getLogger(__name__)
@@ -86,13 +86,15 @@ def __init__(
         context_extractor: ContextExtractor | None = None,
         sampling_rate: float = 1.0,
         instrument_name: str = DEFAULT_INSTRUMENT_NAME,
+        enrich_logger: bool = True,
     ) -> None:
         """Initialize the plugin with an OpenTelemetry tracer provider.
 
         The provided tracer provider is configured with this plugin's
         deterministic ID generator and sampling strategy so spans for a durable
         execution share stable trace and logical operation identifiers.
         """
+        self._enrich_logger = enrich_logger
         self._context_extractor: ContextExtractor = (
             context_extractor or xray_context_extractor
         )
@@ -115,6 +117,24 @@ def __init__(
         self._operation_spans: dict[str | None, Span] = {}
         self._operation_spans_lock = threading.RLock()
 
+    def wrap_logger(self, logger: LoggerInterface) -> LoggerInterface | None:
+        """Wrap the execution logger to inject OTel trace context.
+
+        When enrich_logger is enabled (default), returns an OtelEnrichedLogger
+        that adds trace_id, span_id, and trace_sampled to every log message.
+        Idempotent: returns None if the logger is already an OtelEnrichedLogger.
+
+        Args:
+            logger: The current logger interface from the execution context.
+
+        Returns:
+            An OtelEnrichedLogger wrapping the input, or None if disabled or
+            already wrapped.
+        """
+        if not self._enrich_logger or isinstance(logger, OtelEnrichedLogger):
+            return None
+        return OtelEnrichedLogger(inner=logger, plugin=self)
+
     def _set_span(self, operation_id: str | None, span: Span) -> None:
         """Register the active span for an operation ID."""
         with self._operation_spans_lock:
@@ -130,6 +150,34 @@ def _get_span(self, operation_id: str | None) -> Span | None:
         with self._operation_spans_lock:
             return self._operation_spans.get(operation_id)
 
+    def get_current_span_context(self) -> SpanContext | None:
+        """Return the span context to use for log correlation.
+
+        Resolution order:
+        1. The span attached to the OTel thread-local context. Inside a step or
+           child context this is the active operation span (attached in
+           on_user_function_start), and between operations it is the enclosing
+           operation span (restored in on_user_function_end).
+        2. The invocation span from the plugin registry. This is the path used
+           for top-level handler code: the invocation span is never attached to
+           the worker thread's context, so the registry is the only way to
+           resolve it.
+
+        Returns:
+            A valid SpanContext, or None if no span is active.
+        """
+        span_context = trace.get_current_span().get_span_context()
+        if span_context and span_context.is_valid:
+            return span_context
+
+        invocation_span = self._get_span(None)
+        if invocation_span:
+            invocation_context = invocation_span.get_span_context()
+            if invocation_context and invocation_context.is_valid:
+                return invocation_context
+
+        return None
+
     # ------------------------------------------------------------------
     # Context resolution
     # ------------------------------------------------------------------
@@ -256,7 +304,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
 
         self._start_span(
             operation_id=None,
-            name=f"invocation",
+            name="invocation",
             attributes=self._extract_attributes(info),
         )
 
@@ -413,7 +461,16 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
         if end_timestamp is not None and end_timestamp == info.start_time:
             end_timestamp += datetime.timedelta(microseconds=1)
         self._end_span(info.operation_id, end_timestamp)
-        # We don't call context.detach because the next operation will override it anyway
+        # Restore the enclosing operation span as current so code that runs
+        # after this operation (e.g. between steps in a child context)
+        # correlates to its enclosing operation, not the operation that just
+        # ended. For a top-level operation (parent_id is None) this is the
+        # invocation span; for a nested operation it is the parent context span.
+        parent_span = self._get_span(info.parent_id) or self._get_span(None)
+        if parent_span:
+            context.attach(
+                trace.set_span_in_context(parent_span, self._extracted_context)
+            )
 
     def _extract_attributes(self, info: Any) -> dict[str, str]:
         """Extract durable execution fields as OpenTelemetry span attributes.