From 2b983b9740f9e4910f6faea5d5222593be021fcf Mon Sep 17 00:00:00 2001
From: Ankcorn <thomasankcorn@gmail.com>
Date: Thu, 28 Mar 2024 11:42:22 +0000
Subject: [PATCH] make tracing key _baselime

---
 CHANGELOG.md              | 4 ++++
 src/lambda/parse-event.ts | 6 +++++-
 src/lambda/propation.ts   | 3 ++-
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 526d3c6..558b710 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,10 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [Unreleased]
+
+* make tracing key _baselime for step-functions
+  
 ## [0.5.1] - 2024-03-28
 
 * Improve step function tracing
diff --git a/src/lambda/parse-event.ts b/src/lambda/parse-event.ts
index 275889e..33a3371 100644
--- a/src/lambda/parse-event.ts
+++ b/src/lambda/parse-event.ts
@@ -135,7 +135,11 @@ function detectService(event: any) {
         return "eventbridge";
     }
 
-    if (process.env.BASELIME_TRACE_STEP_FUNCTION === "true" || event.Payload?.baselime?.traceparent || (Array.isArray(event) && event[0]?.Payload?.baselime?.traceparent)) {
+    if (
+        process.env.BASELIME_TRACE_STEP_FUNCTION === "true" ||
+        event.Payload?._baselime?.traceparent || event._baselime?.traceparent ||
+        (Array.isArray(event) && (event[0]?.Payload?.baselime?.traceparent || event[0]?._baselime?.traceparent))
+    ) {
         return "step-function";
     }
     return 'unknown'
diff --git a/src/lambda/propation.ts b/src/lambda/propation.ts
index 36e9700..abae274 100644
--- a/src/lambda/propation.ts
+++ b/src/lambda/propation.ts
@@ -51,9 +51,10 @@ export function extractContext(service: string, event: any): { parent?: Context,
                     }).filter(el => el)
                 }
             }
+            const traceparent = event._baselime?.traceparent || event.Payload?._baselime.traceparent;
 
             return {
-                parent: propagation.extract(context.active(), event?.baselime || event.Payload?.baselime, headerGetter)
+                parent: propagation.extract(context.active(), { traceparent }, headerGetter)
             }
         default:
             return {