diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts index 6a9a9e20..f70865fa 100644 --- a/src/trace/listener.spec.ts +++ b/src/trace/listener.spec.ts @@ -544,4 +544,28 @@ describe("TraceListener", () => { ); }); }); + + it("sets durable function context tags directly on the aws.lambda span", async () => { + const mockSetTag = jest.fn(); + const mockSpan = { setTag: mockSetTag }; + const currentSpanSpy = jest.spyOn(TracerWrapper.prototype, "currentSpan", "get").mockReturnValue(mockSpan); + + try { + const listener = new TraceListener(defaultConfig); + const durableEvent = { + DurableExecutionArn: + "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + }; + await listener.onStartInvocation(durableEvent, context as any); + listener.onEndingInvocation(durableEvent, {}, false); + + expect(mockSetTag).toHaveBeenCalledWith("aws_lambda.durable_function.execution_name", "my-execution"); + expect(mockSetTag).toHaveBeenCalledWith( + "aws_lambda.durable_function.execution_id", + "550e8400-e29b-41d4-a716-446655440004", + ); + } finally { + currentSpanSpy.mockRestore(); + } + }); }); diff --git a/src/trace/listener.ts b/src/trace/listener.ts index d846d7d4..217ba7a8 100644 --- a/src/trace/listener.ts +++ b/src/trace/listener.ts @@ -226,6 +226,14 @@ export class TraceListener { } } } + if (this.durableFunctionContext) { + logDebug("Applying durable function context to the aws.lambda span"); + for (const [key, value] of Object.entries(this.durableFunctionContext)) { + if (value !== undefined) { + this.tracerWrapper.currentSpan.setTag(key, value); + } + } + } let rootSpan = this.inferredSpan; if (!rootSpan) { @@ -337,13 +345,6 @@ export class TraceListener { ...this.stepFunctionContext, }; } - if (this.durableFunctionContext) { - logDebug("Applying durable function context to the aws.lambda span"); - options.tags = { - ...options.tags, - ...this.durableFunctionContext, - }; - } if (this.lambdaSpanParentContext) { options.childOf = this.lambdaSpanParentContext; }