diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index 0e888fef..a4dcd966 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -907,6 +907,43 @@ describe("TraceContextExtractor", () => { expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); }); + + it("returns StepFunctionEventTraceExtractor when event contains legacy lambda StepFunctionContext", () => { + const event = { + Payload: { + Execution: { + Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf", + Input: { + MyInput: "MyValue", + }, + Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be", + RoleArn: + "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03", + StartTime: "2022-12-08T21:08:17.924Z", + }, + State: { + Name: "step-one", + EnteredTime: "2022-12-08T21:08:19.224Z", + RetryCount: 2, + }, + StateMachine: { + Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential", + Name: "my-state-machine", + }, + }, + }; + + const tracerWrapper = new TracerWrapper(); + const traceContextExtractor = new TraceContextExtractor(tracerWrapper, {} as TraceConfig); + + // Mimick TraceContextService.extract initialization + const instance = StepFunctionContextService.instance(event); + traceContextExtractor["stepFunctionContextService"] = instance; + + const extractor = traceContextExtractor["getTraceEventExtractor"](event); + + expect(extractor).toBeInstanceOf(StepFunctionEventTraceExtractor); + }); }); describe("addTraceContexToXray", () => { diff --git a/src/trace/context/extractors/step-function.spec.ts b/src/trace/context/extractors/step-function.spec.ts index 199ad759..37847cf4 100644 --- a/src/trace/context/extractors/step-function.spec.ts +++ b/src/trace/context/extractors/step-function.spec.ts @@ -55,6 +55,22 @@ describe("StepFunctionEventTraceExtractor", () => { expect(traceContext?.source).toBe("event"); }); + it("extracts trace context with valid legacy lambda payload", () => { + // Mimick TraceContextService.extract initialization + StepFunctionContextService.instance({ Payload: payload }); + + const extractor = new StepFunctionEventTraceExtractor(); + + // Payload is sent again for safety in case the instance wasn't previously initialized + const traceContext = extractor.extract({ Payload: payload }); + expect(traceContext).not.toBeNull(); + + expect(traceContext?.toTraceId()).toBe("1139193989631387307"); + expect(traceContext?.toSpanId()).toBe("5892738536804826142"); + expect(traceContext?.sampleMode()).toBe("1"); + expect(traceContext?.source).toBe("event"); + }); + it("returns null when StepFunctionContextService.context is undefined", async () => { const extractor = new StepFunctionEventTraceExtractor(); diff --git a/src/trace/step-function-service.spec.ts b/src/trace/step-function-service.spec.ts index 481a98da..c8a888d0 100644 --- a/src/trace/step-function-service.spec.ts +++ b/src/trace/step-function-service.spec.ts @@ -209,6 +209,21 @@ describe("StepFunctionContextService", () => { expect(spanContext).toBeNull(); }); + + it("returns a SpanContextWrapper when event is from legacy lambda", () => { + const instance = StepFunctionContextService.instance(); + // Force setting event + instance["setContext"]({ Payload: stepFunctionEvent }); + + const spanContext = instance.spanContext; + + expect(spanContext).not.toBeNull(); + + expect(spanContext?.toTraceId()).toBe("1139193989631387307"); + expect(spanContext?.toSpanId()).toBe("5892738536804826142"); + expect(spanContext?.sampleMode()).toBe("1"); + expect(spanContext?.source).toBe("event"); + }); }); describe("deterministicSha256HashToBigIntString", () => { diff --git a/src/trace/step-function-service.ts b/src/trace/step-function-service.ts index d438a78f..5f2e3d9f 100644 --- a/src/trace/step-function-service.ts +++ b/src/trace/step-function-service.ts @@ -41,6 +41,11 @@ export class StepFunctionContextService { // always triggered by the same event. if (typeof event !== "object") return; + // Legacy lambda parsing + if (typeof event.Payload === "object") { + event = event.Payload; + } + // Execution const execution = event.Execution; if (typeof execution !== "object") {