fix: honor upstream trace context in Step Functions Execution.Input._datadog (#830)

Legacy Step Functions invocations were ignoring upstream Datadog trace
headers nested under Execution.Input._datadog and always hashing
Execution.Id instead. Extract shared helper and read _datadog before
falling back to deterministic trace id generation.

APMSVLS-513

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: chemystery09

datadog_lambda/tracing.py

@@ -482,6 +482,21 @@ def _generate_sfn_parent_id(context: dict) -> int:
     )
 
 
+def _extract_dd_trace_id_from_dd_data(dd_data, meta):
+    """
+    Read native Datadog trace headers from a `_datadog` dict.
+    Returns the 64-bit trace_id and populates meta["_dd.p.tid"] when present.
+    Returns None if no usable trace id is found.
+    """
+    if not dd_data or "x-datadog-trace-id" not in dd_data:
+        return None
+    trace_id = int(dd_data.get("x-datadog-trace-id"))
+    high_64_bit_trace_id = _parse_high_64_bits(dd_data.get("x-datadog-tags"))
+    if high_64_bit_trace_id:
+        meta["_dd.p.tid"] = high_64_bit_trace_id
+    return trace_id
+
+
 def _generate_sfn_trace_id(execution_id: str, part: str):
     """
     Take the SHA-256 hash of the execution_id to calculate the trace ID. If the high 64 bits are
@@ -517,10 +532,7 @@ def extract_context_from_step_functions(event, lambda_context):
 
         if event.get("serverless-version") == "v1":
             if "x-datadog-trace-id" in event:  # lambda root
-                trace_id = int(event.get("x-datadog-trace-id"))
-                high_64_bit_trace_id = _parse_high_64_bits(event.get("x-datadog-tags"))
-                if high_64_bit_trace_id:
-                    meta["_dd.p.tid"] = high_64_bit_trace_id
+                trace_id = _extract_dd_trace_id_from_dd_data(event, meta)
             else:  # sfn root
                 root_execution_id = event.get("RootExecutionId")
                 trace_id = _generate_sfn_trace_id(root_execution_id, LOWER_64_BITS)
@@ -530,9 +542,13 @@ def extract_context_from_step_functions(event, lambda_context):
 
             parent_id = _generate_sfn_parent_id(event)
         else:
-            execution_id = event.get("Execution").get("Id")
-            trace_id = _generate_sfn_trace_id(execution_id, LOWER_64_BITS)
-            meta["_dd.p.tid"] = _generate_sfn_trace_id(execution_id, HIGHER_64_BITS)
+            execution = event.get("Execution", {})
+            dd_input = execution.get("Input", {}).get("_datadog")
+            trace_id = _extract_dd_trace_id_from_dd_data(dd_input, meta)
+            if trace_id is None:
+                execution_id = execution.get("Id")
+                trace_id = _generate_sfn_trace_id(execution_id, LOWER_64_BITS)
+                meta["_dd.p.tid"] = _generate_sfn_trace_id(execution_id, HIGHER_64_BITS)
             parent_id = _generate_sfn_parent_id(event)
 
         sampling_priority = SamplingPriority.AUTO_KEEP

tests/test_tracing.py

@@ -751,6 +751,107 @@ def test_step_function_trace_data(self):
             sfn_event, 435175499815315247, 3929055471293792800, "3e7a89d1b7310603"
         )
 
+    @with_trace_propagation_style("datadog")
+    def test_step_function_trace_data_with_input_datadog(self):
+        """Legacy SFN context with upstream trace in Execution.Input._datadog"""
+        sfn_event = {
+            "Execution": {
+                "Id": "arn:aws:states:sa-east-1:425362996713:execution:rstrat-sfn-lambda-jsonpath-demo-dev-state-machine:7d67bd1e-57df-4021-8b52-f55d5cad1169",
+                "Input": {
+                    "hello": "world",
+                    "_datadog": {
+                        "x-datadog-trace-id": "12766418539254701015",
+                        "x-datadog-parent-id": "5497030307431673011",
+                        "x-datadog-sampling-priority": "1",
+                        "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=69f4ed3a00000000",
+                        "traceparent": "00-69f4ed3a00000000b12b6e05a63cdbd7-4c495ff0aa056cb3-01",
+                        "tracestate": "dd=p:4c495ff0aa056cb3;s:1;t.dm:-0;t.tid:69f4ed3a00000000",
+                    },
+                },
+                "StartTime": "2026-05-01T18:13:14.130Z",
+                "Name": "7d67bd1e-57df-4021-8b52-f55d5cad1169",
+                "RoleArn": "arn:aws:iam::425362996713:role/rstrat-sfn-lambda-jsonpat-StepFunctionsExecutionRol-mtfEz4KEzBEE",
+                "RedriveCount": 0,
+            },
+            "State": {
+                "Name": "InvokeDownstreamSecond",
+                "EnteredTime": "2026-05-01T18:13:16.310Z",
+                "RetryCount": 0,
+            },
+            "StateMachine": {
+                "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:rstrat-sfn-lambda-jsonpath-demo-dev-state-machine",
+                "Name": "rstrat-sfn-lambda-jsonpath-demo-dev-state-machine",
+            },
+        }
+        self._test_step_function_trace_data_common(
+            sfn_event, 12766418539254701015, 5718818197702795437, "69f4ed3a00000000"
+        )
+
+    @with_trace_propagation_style("datadog")
+    def test_step_function_trace_data_input_datadog_without_tid(self):
+        """Execution.Input._datadog without _dd.p.tid still preserves upstream trace id"""
+        sfn_event = {
+            "Execution": {
+                "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
+                "Input": {
+                    "_datadog": {
+                        "x-datadog-trace-id": "5821803790426892636",
+                    },
+                },
+                "Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
+                "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
+                "StartTime": "2024-12-04T19:38:04.069Z",
+                "RedriveCount": 0,
+            },
+            "State": {
+                "Name": "Lambda Invoke",
+                "EnteredTime": "2024-12-04T19:38:04.118Z",
+                "RetryCount": 0,
+            },
+            "StateMachine": {
+                "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
+                "Name": "abhinav-activity-state-machine",
+            },
+        }
+        lambda_ctx = get_mock_context()
+        expected_context = Context(
+            trace_id=5821803790426892636,
+            span_id=3929055471293792800,
+            sampling_priority=1,
+            meta={},
+        )
+
+        ctx, source, _ = extract_dd_trace_context(sfn_event, lambda_ctx)
+
+        self.assertEqual(source, "event")
+        self.assertEqual(ctx, expected_context)
+
+    @with_trace_propagation_style("datadog")
+    def test_step_function_trace_data_input_without_datadog(self):
+        """Execution.Input without _datadog falls back to deterministic trace id"""
+        sfn_event = {
+            "Execution": {
+                "Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
+                "Input": {"hello": "world"},
+                "Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
+                "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
+                "StartTime": "2024-12-04T19:38:04.069Z",
+                "RedriveCount": 0,
+            },
+            "State": {
+                "Name": "Lambda Invoke",
+                "EnteredTime": "2024-12-04T19:38:04.118Z",
+                "RetryCount": 0,
+            },
+            "StateMachine": {
+                "Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
+                "Name": "abhinav-activity-state-machine",
+            },
+        }
+        self._test_step_function_trace_data_common(
+            sfn_event, 435175499815315247, 3929055471293792800, "3e7a89d1b7310603"
+        )
+
     @with_trace_propagation_style("datadog")
     def test_step_function_trace_data_retry(self):
         """Test step function trace data extraction with non-zero retry count"""