fix: reset tracecontext to avoid unintentional caching (#676)

* reset tracecontext to avoid unintentional caching

* fix unittest mock

* yarn format

* fix mock in index.spec.ts

* yarn format

Author: joeyzhao2018

src/index.spec.ts

@@ -63,6 +63,10 @@ jest.mock("./trace/trace-context-service", () => {
     get currentTraceHeaders() {
       return mockTraceHeaders;
     }
+
+    reset() {
+      // mocking
+    }
   }
   return {
     ...jest.requireActual("./trace/trace-context-service"),

src/trace/listener.spec.ts

@@ -68,6 +68,9 @@ jest.mock("./trace-context-service", () => {
     get currentTraceContext() {
       return mockSpanContextWrapper;
     }
+    reset() {
+      // mocking
+    }
   }
   return {
     ...jest.requireActual("./trace-context-service"),

src/trace/listener.ts

@@ -279,9 +279,10 @@ export class TraceListener {
       this.injectAuthorizerSpan(result, event?.requestContext?.requestId, finishTime || Date.now());
     }
 
-    // Reset singleton
+    // Reset singletons and trace context
     this.stepFunctionContext = undefined;
     StepFunctionContextService.reset();
+    this.contextService.reset();
   }
 
   public onWrap<T = (...args: any[]) => any>(func: T): T {

src/trace/trace-context-service.spec.ts

@@ -79,4 +79,51 @@ describe("TraceContextService", () => {
     expect(currentTraceContext?.sampleMode()).toBe(1);
     expect(currentTraceContext?.source).toBe("xray");
   });
+
+  it("resets rootTraceContext to prevent caching between invocations", () => {
+    // Initial trace context
+    traceContextService["rootTraceContext"] = {
+      toTraceId: () => "123456",
+      toSpanId: () => "abcdef",
+      sampleMode: () => 1,
+      source: TraceSource.Event,
+      spanContext: spanContext,
+    };
+
+    expect(traceContextService.currentTraceContext).not.toBeNull();
+    expect(traceContextService.traceSource).toBe("event");
+
+    traceContextService.reset();
+
+    expect(traceContextService.currentTraceContext).toBeNull();
+    expect(traceContextService.traceSource).toBeNull();
+  });
+
+  it("automatically resets trace context at the beginning of extract", async () => {
+    // Mock the extractor to return a specific context
+    const mockExtract = jest.fn().mockResolvedValue({
+      toTraceId: () => "newTraceId",
+      toSpanId: () => "newSpanId",
+      sampleMode: () => 1,
+      source: TraceSource.Event,
+      spanContext: {},
+    });
+    traceContextService["traceExtractor"] = { extract: mockExtract } as any;
+
+    // Set up old trace context (simulating previous invocation)
+    traceContextService["rootTraceContext"] = {
+      toTraceId: () => "oldTraceId",
+      toSpanId: () => "oldSpanId",
+      sampleMode: () => 0,
+      source: TraceSource.Xray,
+      spanContext: {},
+    };
+
+    // Extract should reset and set new context
+    const result = await traceContextService.extract({}, {} as any);
+
+    // Verify old context was cleared and new context was set
+    expect(result?.toTraceId()).toBe("newTraceId");
+    expect(traceContextService.traceSource).toBe("event");
+  });
 });

src/trace/trace-context-service.ts

@@ -50,6 +50,9 @@ export class TraceContextService {
   }
 
   async extract(event: any, context: Context): Promise<SpanContextWrapper | null> {
+    // Reset trace context from previous invocation to prevent caching
+    this.rootTraceContext = null;
+
     this.rootTraceContext = await this.traceExtractor?.extract(event, context);
 
     return this.currentTraceContext;
@@ -82,4 +85,8 @@ export class TraceContextService {
   get traceSource() {
     return this.rootTraceContext !== null ? this.rootTraceContext?.source : null;
   }
+
+  reset() {
+    this.rootTraceContext = null;
+  }
 }