fix: ensure aws-xray-propagator maintains trace state

Author: majanjua-amzn

packages/propagator-aws-xray/CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### Bug Fixes
 
+* fix: ensure aws-xray-propagator propagates trace state ([#3092](https://github.com/open-telemetry/opentelemetry-js-contrib/pull/3092))
+
 * **deps:** update all patch versions ([#2948](https://github.com/open-telemetry/opentelemetry-js-contrib/issues/2948)) ([5836d7a](https://github.com/open-telemetry/opentelemetry-js-contrib/commit/5836d7ab3244adef62b715ef22a26b54dba6719b))
 
 ## [2.1.0](https://github.com/open-telemetry/opentelemetry-js-contrib/compare/propagator-aws-xray-v2.0.0...propagator-aws-xray-v2.1.0) (2025-05-02)

packages/propagator-aws-xray/src/AWSXRayPropagator.ts

@@ -81,7 +81,17 @@ export class AWSXRayPropagator implements TextMapPropagator {
     const spanContext = this.getSpanContextFromHeader(carrier, getter);
     if (!isSpanContextValid(spanContext)) return context;
 
-    return trace.setSpan(context, trace.wrapSpanContext(spanContext));
+    // If a previous propagator already set the trace state, ensure it's propagated
+    const existingSpan = trace.getSpan(context);
+    const existingTraceState = existingSpan?.spanContext()?.traceState;
+    const finalSpanContext: SpanContext = existingTraceState
+      ? {
+          ...spanContext,
+          traceState: existingTraceState,
+        }
+      : spanContext;
+
+    return trace.setSpan(context, trace.wrapSpanContext(finalSpanContext));
   }
 
   fields(): string[] {

packages/propagator-aws-xray/test/AWSXRayPropagator.test.ts

@@ -79,9 +79,8 @@ describe('AWSXRayPropagator', () => {
       );
     });
 
-    it('should inject with TraceState', () => {
-      const traceState = new TraceState();
-      traceState.set('foo', 'bar');
+    it('should not inject with TraceState', () => {
+      const traceState = new TraceState().set('foo', 'bar').set('key', 'val');
       const spanContext: SpanContext = {
         traceId: TRACE_ID,
         spanId: SPAN_ID,
@@ -94,7 +93,8 @@ describe('AWSXRayPropagator', () => {
         defaultTextMapSetter
       );
 
-      // TODO: assert trace state when the propagator supports it
+      // Rely on W3CTraceContextPropagator to propagate trace state
+      // As such, it's expected that trace state is not populated here
       assert.deepStrictEqual(
         carrier[AWSXRAY_TRACE_ID_HEADER],
         'Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=1'
@@ -345,6 +345,36 @@ describe('AWSXRayPropagator', () => {
       });
     });
 
+    it('should extract with TraceState previously propagated', () => {
+      carrier[AWSXRAY_TRACE_ID_HEADER] =
+        'Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=1';
+
+      const spanContextWithTraceState: SpanContext = {
+        traceId: 'existing-trace-id',
+        spanId: 'existing-span-id',
+        traceFlags: TraceFlags.SAMPLED,
+        traceState: new TraceState().set('foo', 'bar').set('from', 'context'),
+      };
+      const incomingContext = trace.setSpan(
+        ROOT_CONTEXT,
+        trace.wrapSpanContext(spanContextWithTraceState)
+      );
+
+      const extractedSpanContext = trace
+        .getSpan(
+          xrayPropagator.extract(incomingContext, carrier, defaultTextMapGetter)
+        )
+        ?.spanContext();
+
+      assert.deepStrictEqual(extractedSpanContext, {
+        traceId: TRACE_ID,
+        spanId: SPAN_ID,
+        isRemote: true,
+        traceFlags: TraceFlags.SAMPLED,
+        traceState: new TraceState().set('foo', 'bar').set('from', 'context'),
+      });
+    });
+
     describe('.fields()', () => {
       it('should return a field with AWS X-Ray Trace ID header', () => {
         const expectedField = xrayPropagator.fields();