From f1854ad46bf7ca057ac07e7e45ded4337a551f86 Mon Sep 17 00:00:00 2001
From: majanjua-amzn <majanjua@amazon.com>
Date: Tue, 14 Oct 2025 09:01:21 -0700
Subject: [PATCH] fix: ensure aws-xray-propagator maintains trace state

---
 .../src/AWSXRayPropagator.ts                  |  6 +++
 .../test/AWSXRayPropagator.test.ts            | 38 +++++++++++++++++--
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/packages/propagator-aws-xray/src/AWSXRayPropagator.ts b/packages/propagator-aws-xray/src/AWSXRayPropagator.ts
index 06f535ad93..a0def93418 100644
--- a/packages/propagator-aws-xray/src/AWSXRayPropagator.ts
+++ b/packages/propagator-aws-xray/src/AWSXRayPropagator.ts
@@ -81,6 +81,12 @@ export class AWSXRayPropagator implements TextMapPropagator {
     const spanContext = this.getSpanContextFromHeader(carrier, getter);
     if (!isSpanContextValid(spanContext)) return context;
 
+    // If a previous propagator already set the trace state, ensure it's propagated
+    const existingSpan = trace.getSpan(context);
+    const existingTraceState = existingSpan?.spanContext()?.traceState;
+    if (existingTraceState) {
+      spanContext.traceState = existingTraceState
+    }
     return trace.setSpan(context, trace.wrapSpanContext(spanContext));
   }
 
diff --git a/packages/propagator-aws-xray/test/AWSXRayPropagator.test.ts b/packages/propagator-aws-xray/test/AWSXRayPropagator.test.ts
index 4a6ab8db1c..60d289d6f0 100644
--- a/packages/propagator-aws-xray/test/AWSXRayPropagator.test.ts
+++ b/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();