From 2016cd318f5ced90c52eb2f97838b4de20452785 Mon Sep 17 00:00:00 2001 From: wangzlei Date: Mon, 20 Apr 2026 16:59:06 -0700 Subject: [PATCH 1/3] fix(lambda): align Python context propagation with JS implementation The custom_event_context_extractor was always using the X-Ray env var as parent context because _X_AMZN_TRACE_ID is always set and valid when active tracing is enabled, causing upstream W3C trace context (traceparent) from event headers to be ignored and breaking trace continuity. Align with the JS implementation: inject the X-Ray env var into the event headers and delegate to the global composite propagator, which respects propagator priority ordering (xray, then tracecontext). --- lambda-layer/src/otel_wrapper.py | 28 +-- .../src/tests/test_lambda_instrumentation.py | 190 +++++++++++++----- 2 files changed, 158 insertions(+), 60 deletions(-) diff --git a/lambda-layer/src/otel_wrapper.py b/lambda-layer/src/otel_wrapper.py index 8ea565e49..9017dbe3c 100644 --- a/lambda-layer/src/otel_wrapper.py +++ b/lambda-layer/src/otel_wrapper.py @@ -41,7 +41,6 @@ from opentelemetry.context import Context from opentelemetry.instrumentation.aws_lambda import _X_AMZN_TRACE_ID, AwsLambdaInstrumentor from opentelemetry.propagate import get_global_textmap -from opentelemetry.propagators.aws import AwsXRayPropagator from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_HEADER_KEY from opentelemetry.trace import get_current_span @@ -57,21 +56,24 @@ class HandlerError(Exception): def custom_event_context_extractor(lambda_event: Any) -> Context: xray_env_var = os.environ.get(_X_AMZN_TRACE_ID) - lambda_trace_context = AwsXRayPropagator().extract({TRACE_HEADER_KEY: xray_env_var}) - parent_span_context = get_current_span(lambda_trace_context).get_span_context() - if parent_span_context is None or not parent_span_context.is_valid: + try: + headers = lambda_event["headers"] + except (TypeError, KeyError): headers = None - try: - headers = lambda_event["headers"] - except (TypeError, KeyError): - pass - if not isinstance(headers, dict): - headers = {} + if not isinstance(headers, dict): + headers = {} + else: + headers = {k: v for k, v in headers.items()} - return get_global_textmap().extract(headers) + if xray_env_var: + headers = {k: v for k, v in headers.items() if k.lower() != TRACE_HEADER_KEY.lower()} + headers[TRACE_HEADER_KEY] = xray_env_var - return lambda_trace_context + extracted_context = get_global_textmap().extract(headers) + if get_current_span(extracted_context).get_span_context(): + return extracted_context + return Context() AwsLambdaInstrumentor().instrument(event_context_extractor=custom_event_context_extractor) @@ -82,7 +84,7 @@ def custom_event_context_extractor(lambda_event: Any) -> Context: raise HandlerError("ORIG_HANDLER is not defined.") try: - (mod_name, handler_name) = path.rsplit(".", 1) + mod_name, handler_name = path.rsplit(".", 1) except ValueError as e: raise HandlerError("Bad path '{}' for ORIG_HANDLER: {}".format(path, str(e))) diff --git a/lambda-layer/src/tests/test_lambda_instrumentation.py b/lambda-layer/src/tests/test_lambda_instrumentation.py index 023f28a2e..3816aff17 100644 --- a/lambda-layer/src/tests/test_lambda_instrumentation.py +++ b/lambda-layer/src/tests/test_lambda_instrumentation.py @@ -25,9 +25,11 @@ from shutil import which from unittest import mock -from opentelemetry.environment_variables import OTEL_PROPAGATORS from opentelemetry.instrumentation.aws_lambda import _HANDLER, _X_AMZN_TRACE_ID, ORIG_HANDLER, AwsLambdaInstrumentor +from opentelemetry.propagate import get_global_textmap, set_global_textmap +from opentelemetry.propagators.aws import AwsXRayPropagator from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_ID_FIRST_PART_LENGTH, TRACE_ID_VERSION +from opentelemetry.propagators.composite import CompositePropagator from opentelemetry.semconv.resource import ResourceAttributes from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.test.test_base import TestBase @@ -196,6 +198,9 @@ def tearDownClass(cls): ) def test_active_tracing(self): + original_propagator = get_global_textmap() + set_global_textmap(AwsXRayPropagator()) + test_env_patch = mock.patch.dict( "os.environ", { @@ -206,71 +211,162 @@ def test_active_tracing(self): ) test_env_patch.start() - mock_execute_lambda() - - spans = self.memory_exporter.get_finished_spans() + try: + mock_execute_lambda() + + spans = self.memory_exporter.get_finished_spans() + + assert spans + + self.assertEqual(len(spans), 1) + span = spans[0] + self.assertEqual(span.name, os.environ[ORIG_HANDLER]) + self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID) + self.assertEqual(span.kind, SpanKind.SERVER) + self.assertSpanHasAttributes( + span, + { + ResourceAttributes.CLOUD_RESOURCE_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn, + SpanAttributes.FAAS_INVOCATION_ID: MOCK_LAMBDA_CONTEXT.aws_request_id, + }, + ) + + parent_context = span.parent + self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id) + self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID) + self.assertTrue(parent_context.is_remote) + finally: + test_env_patch.stop() + set_global_textmap(original_propagator) - assert spans + def test_parent_context_from_lambda_event(self): + original_propagator = get_global_textmap() + set_global_textmap(TraceContextTextMapPropagator()) - self.assertEqual(len(spans), 1) - span = spans[0] - self.assertEqual(span.name, os.environ[ORIG_HANDLER]) - self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID) - self.assertEqual(span.kind, SpanKind.SERVER) - self.assertSpanHasAttributes( - span, + test_env_patch = mock.patch.dict( + "os.environ", { - ResourceAttributes.CLOUD_RESOURCE_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn, - SpanAttributes.FAAS_INVOCATION_ID: MOCK_LAMBDA_CONTEXT.aws_request_id, + **os.environ, + # NOT Active Tracing + _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH, }, ) + test_env_patch.start() - parent_context = span.parent - self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id) - self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID) - self.assertTrue(parent_context.is_remote) - - test_env_patch.stop() + try: + trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2" + mock_execute_lambda( + { + "headers": { + TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED, + TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header, + } + } + ) + + spans = self.memory_exporter.get_finished_spans() + + assert spans + + self.assertEqual(len(spans), 1) + span = spans[0] + self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID) + + parent_context = span.parent + self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id) + self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID) + self.assertEqual(len(parent_context.trace_state), 3) + self.assertEqual( + parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY), + MOCK_W3C_TRACE_STATE_VALUE, + ) + self.assertTrue(parent_context.is_remote) + finally: + test_env_patch.stop() + set_global_textmap(original_propagator) + + def test_xray_ignored_when_propagator_does_not_include_xray(self): + """When X-Ray active tracing is on but the user only configures + tracecontext propagator (no xray), the X-Ray env var is injected + into headers but the propagator cannot parse it. The W3C context + from event headers should be used instead.""" + original_propagator = get_global_textmap() + set_global_textmap(TraceContextTextMapPropagator()) - def test_parent_context_from_lambda_event(self): test_env_patch = mock.patch.dict( "os.environ", { **os.environ, - # NOT Active Tracing - _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH, - # NOT using the X-Ray Propagator - OTEL_PROPAGATORS: "tracecontext", + _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED, }, ) test_env_patch.start() - trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2" - mock_execute_lambda( - { - "headers": { - TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED, - TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header, + try: + mock_execute_lambda( + { + "headers": { + TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED, + } } - } - ) + ) + + spans = self.memory_exporter.get_finished_spans() - spans = self.memory_exporter.get_finished_spans() + assert spans + self.assertEqual(len(spans), 1) + span = spans[0] - assert spans + self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID) - self.assertEqual(len(spans), 1) - span = spans[0] - self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID) + parent_context = span.parent + self.assertEqual(parent_context.trace_id, MOCK_W3C_TRACE_ID) + self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID) + self.assertTrue(parent_context.is_remote) + finally: + test_env_patch.stop() + set_global_textmap(original_propagator) - parent_context = span.parent - self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id) - self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID) - self.assertEqual(len(parent_context.trace_state), 3) - self.assertEqual( - parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY), - MOCK_W3C_TRACE_STATE_VALUE, + def test_w3c_takes_precedence_over_xray_when_both_present(self): + """When both X-Ray active tracing and W3C traceparent headers are + present, the W3C context should be used as the parent because + tracecontext comes after xray in the composite propagator — the + last propagator to extract a valid context wins.""" + original_propagator = get_global_textmap() + set_global_textmap(CompositePropagator([AwsXRayPropagator(), TraceContextTextMapPropagator()])) + + test_env_patch = mock.patch.dict( + "os.environ", + { + **os.environ, + _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED, + }, ) - self.assertTrue(parent_context.is_remote) + test_env_patch.start() + + try: + trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2" + mock_execute_lambda( + { + "headers": { + TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED, + TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header, + } + } + ) + + spans = self.memory_exporter.get_finished_spans() + + assert spans + self.assertEqual(len(spans), 1) + span = spans[0] + + self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID) - test_env_patch.stop() + parent_context = span.parent + self.assertEqual(parent_context.trace_id, MOCK_W3C_TRACE_ID) + self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID) + self.assertTrue(parent_context.is_remote) + finally: + test_env_patch.stop() + set_global_textmap(original_propagator) From e1478c640d55dab8bfc152f8e0bd018c28a2cf86 Mon Sep 17 00:00:00 2001 From: wangzlei Date: Mon, 20 Apr 2026 17:11:48 -0700 Subject: [PATCH 2/3] docs: add CHANGELOG entry for lambda context propagation fix --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b1dcdd0d..bab6f1153 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ If your change does not need a CHANGELOG entry, add the "skip changelog" label t ## Unreleased +- fix(lambda-layer): align context propagation with JS — delegate to global propagator so W3C traceparent is no longer ignored when X-Ray active tracing is enabled + ([#727](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/727)) - feat: support environment-configured endpoint visibility for HTTP operation names ([#718](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/718)) - fix(lambda-layer): Disable all agentic instrumentation in Lambda by default From b33ebbb5a11b3f9498fca97d15f8fe5162afda58 Mon Sep 17 00:00:00 2001 From: wangzlei Date: Wed, 22 Apr 2026 21:15:34 -0700 Subject: [PATCH 3/3] refactor: use headers.copy() instead of dict comprehension for shallow copy --- lambda-layer/src/otel_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lambda-layer/src/otel_wrapper.py b/lambda-layer/src/otel_wrapper.py index 9017dbe3c..67fe2ef5c 100644 --- a/lambda-layer/src/otel_wrapper.py +++ b/lambda-layer/src/otel_wrapper.py @@ -64,7 +64,7 @@ def custom_event_context_extractor(lambda_event: Any) -> Context: if not isinstance(headers, dict): headers = {} else: - headers = {k: v for k, v in headers.items()} + headers = headers.copy() if xray_env_var: headers = {k: v for k, v in headers.items() if k.lower() != TRACE_HEADER_KEY.lower()}