Add durable_function_first_invocation tag to aws.lambda spans

Extends extract_durable_function_tags() to set
durable_function_first_invocation ("true"/"false") by checking
len(InitialExecutionState.Operations) <= 1, mirroring the SDK's
own replay detection logic (ReplayStatus.REPLAY when len > 1).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: lym953

datadog_lambda/durable.py

@@ -24,7 +24,7 @@ def _parse_durable_execution_arn(arn):
     return execution_name, execution_id
 
 
-def extract_durable_function_tags(event, state=None):
+def extract_durable_function_tags(event):
     """
     Extracts durable function tags from the Lambda event payload.
     Returns a dict with durable function tags, or an empty dict if the event
@@ -43,15 +43,10 @@ def extract_durable_function_tags(event, state=None):
         return {}
 
     execution_name, execution_id = parsed
-    try:
-        first_invocation = str(not state.is_replaying()).lower()
-    except Exception:
-        logger.debug("Failed to compute durable_function_first_invocation")
-        first_invocation = None
-    tags = {
+    operations = event.get("InitialExecutionState", {}).get("Operations", [])
+    is_first_invocation = len(operations) <= 1
+    return {
         "durable_function_execution_name": execution_name,
         "durable_function_execution_id": execution_id,
+        "durable_function_first_invocation": str(is_first_invocation).lower(),
     }
-    if first_invocation is not None:
-        tags["durable_function_first_invocation"] = first_invocation
-    return tags

datadog_lambda/wrapper.py

@@ -153,7 +153,7 @@ def __init__(self, func):
             if config.trace_extractor:
                 extractor_parts = config.trace_extractor.rsplit(".", 1)
                 if len(extractor_parts) == 2:
-                    mod_name, extractor_name = extractor_parts
+                    (mod_name, extractor_name) = extractor_parts
                     modified_extractor_name = modify_module_name(mod_name)
                     extractor_module = import_module(modified_extractor_name)
                     self.trace_extractor = getattr(extractor_module, extractor_name)
@@ -244,9 +244,7 @@ def _before(self, event, context):
                 submit_invocations_metric(context)
 
             self.trigger_tags = extract_trigger_tags(event, context)
-            self.durable_function_tags = extract_durable_function_tags(
-                event, context.state
-            )
+            self.durable_function_tags = extract_durable_function_tags(event)
             # Extract Datadog trace context and source from incoming requests
             dd_context, trace_context_source, event_source = extract_dd_trace_context(
                 event,

tests/test_durable.py

@@ -3,7 +3,6 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 import unittest
-from unittest.mock import MagicMock
 
 from datadog_lambda.durable import (
     _parse_durable_execution_arn,
@@ -46,31 +45,31 @@ def test_works_with_numeric_version_qualifier(self):
 
 
 class TestExtractDurableFunctionTags(unittest.TestCase):
-    def test_extracts_tags_from_event_with_durable_execution_arn(self):
+    def test_sets_first_invocation_true_when_only_execution_operation(self):
+        # One operation (the EXECUTION op itself) → not replaying → first invocation
         event = {
             "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
             "CheckpointToken": "some-token",
-            "InitialExecutionState": {"Operations": []},
+            "InitialExecutionState": {"Operations": [{"OperationType": "EXECUTION"}]},
         }
-        state = MagicMock()
-        state.is_replaying.return_value = True
-        result = extract_durable_function_tags(event, state)
+        result = extract_durable_function_tags(event)
         self.assertEqual(
             result,
             {
                 "durable_function_execution_name": "my-execution",
                 "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004",
-                "durable_function_first_invocation": "false",
+                "durable_function_first_invocation": "true",
             },
         )
 
-    def test_sets_first_invocation_true_when_not_replaying(self):
+    def test_sets_first_invocation_true_when_no_operations(self):
+        # Empty operations list → not replaying → first invocation
         event = {
             "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
+            "CheckpointToken": "some-token",
+            "InitialExecutionState": {"Operations": []},
         }
-        state = MagicMock()
-        state.is_replaying.return_value = False
-        result = extract_durable_function_tags(event, state)
+        result = extract_durable_function_tags(event)
         self.assertEqual(
             result,
             {
@@ -80,13 +79,19 @@ def test_sets_first_invocation_true_when_not_replaying(self):
             },
         )
 
-    def test_sets_first_invocation_false_when_replaying(self):
+    def test_sets_first_invocation_false_when_multiple_operations(self):
+        # More than one operation → replaying → not first invocation
         event = {
             "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
+            "CheckpointToken": "some-token",
+            "InitialExecutionState": {
+                "Operations": [
+                    {"OperationType": "EXECUTION"},
+                    {"OperationType": "STEP"},
+                ]
+            },
         }
-        state = MagicMock()
-        state.is_replaying.return_value = True
-        result = extract_durable_function_tags(event, state)
+        result = extract_durable_function_tags(event)
         self.assertEqual(
             result,
             {
@@ -96,6 +101,22 @@ def test_sets_first_invocation_false_when_replaying(self):
             },
         )
 
+    def test_sets_first_invocation_true_when_initial_execution_state_absent(self):
+        # No InitialExecutionState key → treated as empty → first invocation
+        event = {
+            "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
+            "CheckpointToken": "some-token",
+        }
+        result = extract_durable_function_tags(event)
+        self.assertEqual(
+            result,
+            {
+                "durable_function_execution_name": "my-execution",
+                "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004",
+                "durable_function_first_invocation": "true",
+            },
+        )
+
     def test_returns_empty_dict_for_regular_lambda_event(self):
         event = {
             "body": '{"key": "value"}',