feat: add w3c headers

.changeset/silly-owls-travel.md

@@ -0,0 +1,5 @@
+---
+"@aws/lambda-invoke-store": minor
+---
+
+add support for traceparent, tracestate and baggage

README.md

@@ -8,7 +8,7 @@ concurrent executions.
 ## Features
 
 - **Invocation Isolation**: Safely store and retrieve data within a single Lambda invocation.
-- **Protected Lambda Context**: Built-in protection for Lambda execution metadata (requestId, [traceId](https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-traces))
+- **Protected Lambda Context**: Built-in protection for Lambda execution metadata (requestId, [traceId](https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-traces), [W3C Trace Context](https://www.w3.org/TR/trace-context/), [W3C Baggage](https://www.w3.org/TR/baggage/))
 - **Custom Data Storage**: Store any custom data within the invocation context
 - **Async/Await Support**: Full support for asynchronous operations with context preservation
 - **Type Safety**: Complete TypeScript type definitions
@@ -123,6 +123,30 @@ Convenience method to get the current [X-Ray trace ID](https://docs.aws.amazon.c
 const traceId = invokeStore.getXRayTraceId(); // Returns undefined if not set or outside context
 ```
 
+### invokeStore.getTraceparent()
+
+Convenience method to get the [W3C `traceparent`](https://www.w3.org/TR/trace-context/#traceparent-header) header value.
+
+```typescript
+const traceparent = invokeStore.getTraceparent(); // e.g. "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+```
+
+### invokeStore.getTracestate()
+
+Convenience method to get the [W3C `tracestate`](https://www.w3.org/TR/trace-context/#tracestate-header) header value.
+
+```typescript
+const tracestate = invokeStore.getTracestate(); // e.g. "congo=t61rcWkgMzE,rojo=00f067aa0ba902b7"
+```
+
+### invokeStore.getBaggage()
+
+Convenience method to get the [W3C Baggage](https://www.w3.org/TR/baggage/) header value.
+
+```typescript
+const baggage = invokeStore.getBaggage(); // e.g. "userId=alice,serverNode=DF28"
+```
+
 ### invokeStore.hasContext()
 
 Checks if code is currently running within an invoke context.

src/invoke-store.global.spec.ts

@@ -120,6 +120,13 @@ describe.each([
         getRequestId: vi.fn().mockReturnValue("mock-request-id"),
         getXRayTraceId: vi.fn(),
         getTenantId: vi.fn().mockReturnValue("my-test-tenant-id"),
+        getTraceparent: vi
+          .fn()
+          .mockReturnValue(
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+          ),
+        getTracestate: vi.fn().mockReturnValue("congo=t61rcWkgMzE"),
+        getBaggage: vi.fn().mockReturnValue("userId=alice,serverNode=DF28"),
         hasContext: vi.fn(),
       };
 
@@ -136,6 +143,13 @@ describe.each([
       expect(awaitedReimportedStore).toBe(mockInstance);
       expect(awaitedReimportedStore.getRequestId()).toBe("mock-request-id");
       expect(awaitedReimportedStore.getTenantId()).toBe("my-test-tenant-id");
+      expect(awaitedReimportedStore.getTraceparent()).toBe(
+        "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+      );
+      expect(awaitedReimportedStore.getTracestate()).toBe("congo=t61rcWkgMzE");
+      expect(awaitedReimportedStore.getBaggage()).toBe(
+        "userId=alice,serverNode=DF28",
+      );
     });
   });
 

src/invoke-store.spec.ts

@@ -43,6 +43,89 @@ describe.each([
       });
     });
 
+    describe("getTraceparent, getTracestate, and getBaggage", () => {
+      it("should return w3c tracing headers when set in context", async () => {
+        // WHEN
+        const result = await invokeStore.run(
+          {
+            [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id",
+            [InvokeStoreBase.PROTECTED_KEYS.TRACEPARENT]:
+              "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+            [InvokeStoreBase.PROTECTED_KEYS.TRACESTATE]: "congo=t61rcWkgMzE",
+            [InvokeStoreBase.PROTECTED_KEYS.BAGGAGE]:
+              "userId=alice,serverNode=DF28",
+          },
+          () => {
+            return {
+              traceparent: invokeStore.getTraceparent(),
+              tracestate: invokeStore.getTracestate(),
+              baggage: invokeStore.getBaggage(),
+            };
+          },
+        );
+
+        // THEN
+        expect(result.traceparent).toBe(
+          "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+        );
+        expect(result.tracestate).toBe("congo=t61rcWkgMzE");
+        expect(result.baggage).toBe("userId=alice,serverNode=DF28");
+      });
+
+      it("should return undefined when w3c headers are not in context", async () => {
+        // WHEN
+        const result = await invokeStore.run(
+          {
+            [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id",
+          },
+          () => {
+            return {
+              traceparent: invokeStore.getTraceparent(),
+              tracestate: invokeStore.getTracestate(),
+              baggage: invokeStore.getBaggage(),
+            };
+          },
+        );
+
+        // THEN
+        expect(result.traceparent).toBeUndefined();
+        expect(result.tracestate).toBeUndefined();
+        expect(result.baggage).toBeUndefined();
+      });
+
+      it("should prevent modifying w3c tracing protected fields", async () => {
+        // WHEN & THEN
+        await invokeStore.run(
+          {
+            [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id",
+            [InvokeStoreBase.PROTECTED_KEYS.TRACEPARENT]: "original",
+          },
+          () => {
+            expect(() => {
+              invokeStore.set(
+                InvokeStoreBase.PROTECTED_KEYS.TRACEPARENT,
+                "modified",
+              );
+            }).toThrow(/Cannot modify protected Lambda context field/);
+
+            expect(() => {
+              invokeStore.set(
+                InvokeStoreBase.PROTECTED_KEYS.TRACESTATE,
+                "modified",
+              );
+            }).toThrow(/Cannot modify protected Lambda context field/);
+
+            expect(() => {
+              invokeStore.set(
+                InvokeStoreBase.PROTECTED_KEYS.BAGGAGE,
+                "modified",
+              );
+            }).toThrow(/Cannot modify protected Lambda context field/);
+          },
+        );
+      });
+    });
+
     describe("custom properties", () => {
       it("should allow setting and getting custom properties", async () => {
         // WHEN

src/invoke-store.ts

@@ -14,6 +14,9 @@ const PROTECTED_KEYS = {
   REQUEST_ID: Symbol.for("_AWS_LAMBDA_REQUEST_ID"),
   X_RAY_TRACE_ID: Symbol.for("_AWS_LAMBDA_X_RAY_TRACE_ID"),
   TENANT_ID: Symbol.for("_AWS_LAMBDA_TENANT_ID"),
+  TRACEPARENT: Symbol.for("_AWS_LAMBDA_TRACEPARENT"),
+  TRACESTATE: Symbol.for("_AWS_LAMBDA_TRACESTATE"),
+  BAGGAGE: Symbol.for("_AWS_LAMBDA_BAGGAGE"),
 } as const;
 
 const NO_GLOBAL_AWS_LAMBDA = ["true", "1"].includes(
@@ -57,6 +60,18 @@ export abstract class InvokeStoreBase {
   getTenantId(): string | undefined {
     return this.get<string>(PROTECTED_KEYS.TENANT_ID);
   }
+
+  getTraceparent(): string | undefined {
+    return this.get<string>(PROTECTED_KEYS.TRACEPARENT);
+  }
+
+  getTracestate(): string | undefined {
+    return this.get<string>(PROTECTED_KEYS.TRACESTATE);
+  }
+
+  getBaggage(): string | undefined {
+    return this.get<string>(PROTECTED_KEYS.BAGGAGE);
+  }
 }
 
 /**