From 390b33a1356691e3ab64cd40a6ecec372a54112e Mon Sep 17 00:00:00 2001
From: Ankcorn <thomasankcorn@gmail.com>
Date: Thu, 28 Mar 2024 14:19:40 +0000
Subject: [PATCH] fix: handle undefined traceparent

---
 CHANGELOG.md                                              | 3 +++
 examples/state-machine/packages/functions/src/lambda.ts   | 2 +-
 examples/state-machine/packages/functions/src/task-two.ts | 7 +++----
 src/lambda/propation.ts                                   | 4 ++--
 4 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c210646..e395c31 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,9 @@ 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]
+
+* Handle undefined traceparent for step function propagation
 
 ## [0.5.2] - 2024-03-28
 
diff --git a/examples/state-machine/packages/functions/src/lambda.ts b/examples/state-machine/packages/functions/src/lambda.ts
index 9ff800f..5c9c39a 100644
--- a/examples/state-machine/packages/functions/src/lambda.ts
+++ b/examples/state-machine/packages/functions/src/lambda.ts
@@ -1,4 +1,4 @@
-import { withOpenTelemetry } from "./tracing";
+// import { withOpenTelemetry } from "./tracing";
 import { ApiHandler } from "sst/node/api";
 
 export const handler = withOpenTelemetry(ApiHandler(async (_evt) => {
diff --git a/examples/state-machine/packages/functions/src/task-two.ts b/examples/state-machine/packages/functions/src/task-two.ts
index a01bdc6..86e5e69 100644
--- a/examples/state-machine/packages/functions/src/task-two.ts
+++ b/examples/state-machine/packages/functions/src/task-two.ts
@@ -1,9 +1,8 @@
-import { withOpenTelemetry } from "./tracing";
+import { ApiHandler } from "sst/node/api";
 
-export const handler = withOpenTelemetry(async (_evt) => {
-  console.log(_evt)
+export const handler = ApiHandler(async (_evt) => {
   return {
     statusCode: 200,
     body: `Hello world. The time is ${new Date().toISOString()}`,
   };
-});
+});
\ No newline at end of file
diff --git a/src/lambda/propation.ts b/src/lambda/propation.ts
index abae274..d1db351 100644
--- a/src/lambda/propation.ts
+++ b/src/lambda/propation.ts
@@ -37,7 +37,7 @@ export function extractContext(service: string, event: any): { parent?: Context,
             if (Array.isArray(event)) {
                 return {
                     links: event.map((parent) => {
-                        const traceparent = parent._baselime?.traceparent || parent.Payload?._baselime.traceparent;
+                        const traceparent = parent._baselime?.traceparent || parent.Payload?._baselime?.traceparent;
                         if (!traceparent) {
                             return
                         }
@@ -51,7 +51,7 @@ export function extractContext(service: string, event: any): { parent?: Context,
                     }).filter(el => el)
                 }
             }
-            const traceparent = event._baselime?.traceparent || event.Payload?._baselime.traceparent;
+            const traceparent = event._baselime?.traceparent || event.Payload?._baselime?.traceparent;
 
             return {
                 parent: propagation.extract(context.active(), { traceparent }, headerGetter)