fix(cold_start): capture trace context before finishing spans

Move trace context capture BEFORE span.finish() so ColdStartTracer
creates spans with the correct trace_id.

Previously, tracer.current_trace_context() was called AFTER span.finish(),
which returned a stale context, causing cold start spans (aws.lambda.load
and import spans) to have a different trace_id than the main aws.lambda span.

Added unit test assertions to verify trace_id consistency.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: jchrostek-dd

datadog_lambda/wrapper.py

@@ -305,6 +305,16 @@ def _after(self, event, context):
 
             status_code = extract_http_status_code_tag(self.trigger_tags, self.response)
 
+            # Skip creating cold start spans in managed instances mode
+            # In managed instances, the tracer library handles cold start independently
+            should_trace_cold_start = (
+                config.cold_start_tracing
+                and is_new_sandbox()
+                and not is_managed_instances_mode()
+            )
+            if should_trace_cold_start:
+                trace_ctx = tracer.current_trace_context()
+
             if self.span:
                 if config.appsec_enabled and not self.blocking_response:
                     asm_start_response(
@@ -342,15 +352,6 @@ def _after(self, event, context):
                 create_dd_dummy_metadata_subsegment(
                     self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
                 )
-            # Skip creating cold start spans in managed instances mode
-            # In managed instances, the tracer library handles cold start independently
-            should_trace_cold_start = (
-                config.cold_start_tracing
-                and is_new_sandbox()
-                and not is_managed_instances_mode()
-            )
-            if should_trace_cold_start:
-                trace_ctx = tracer.current_trace_context()
 
             if self.inferred_span:
                 if status_code:

tests/test_cold_start.py

@@ -310,13 +310,18 @@ def handler(event, context):
     handler.cold_start_tracing = True
     handler({}, lambda_context)
 
-    function_span = import_span = None
+    function_span = import_span = load_span = None
     for span in spans:
         if span.resource == "tabnanny":
             import_span = span
         elif span.name == "aws.lambda":
             function_span = span
+        elif span.name == "aws.lambda.load":
+            load_span = span
 
     assert function_span is not None
     assert import_span is not None
     assert import_span.parent_id == function_span.span_id
+    assert import_span.trace_id == function_span.trace_id
+    assert load_span is not None
+    assert load_span.trace_id == function_span.trace_id