Improve dynamic sampling context and baggage consistency

Author: onurtemizkan

dev-packages/e2e-tests/test-applications/remix-server-timing/tests/server-timing-trace-propagation.test.ts

@@ -55,14 +55,16 @@ test('Propagates trace context from Server-Timing header to client pageload (wit
   expect(serverTimingHeader).toContain('sentry-trace');
   expect(serverTimingHeader).toContain('baggage');
 
-  // Extract the trace ID from the Server-Timing header
+  // Extract the trace ID and span ID from the Server-Timing header
   const sentryTraceMatch = serverTimingHeader?.match(/sentry-trace;desc="([^"]+)"/);
   expect(sentryTraceMatch).toBeTruthy();
 
   const serverTimingTraceValue = sentryTraceMatch?.[1];
-  const serverTimingTraceId = serverTimingTraceValue?.split('-')[0];
+  const [serverTimingTraceId, serverTimingSpanId] = serverTimingTraceValue?.split('-') || [];
   expect(serverTimingTraceId).toBeDefined();
   expect(serverTimingTraceId).toHaveLength(32);
+  expect(serverTimingSpanId).toBeDefined();
+  expect(serverTimingSpanId).toHaveLength(16);
 
   const pageloadTransaction = await pageLoadTransactionPromise;
 
@@ -76,8 +78,9 @@ test('Propagates trace context from Server-Timing header to client pageload (wit
   // This proves that the client correctly read and used the Server-Timing header for trace propagation
   expect(pageLoadTraceId).toEqual(serverTimingTraceId);
 
-  // Verify that the pageload has a parent span (indicating it's continuing a trace, not starting a new one)
-  expect(pageLoadParentSpanId).toBeDefined();
+  // Verify that the pageload's parent_span_id matches the span_id from Server-Timing
+  // This proves the client is continuing from the exact server span, not just using the trace_id
+  expect(pageLoadParentSpanId).toEqual(serverTimingSpanId);
 });
 
 test('Propagates trace context via Server-Timing for parameterized routes', async ({ page }) => {
@@ -124,7 +127,7 @@ test('Propagates trace context via Server-Timing for parameterized routes', asyn
   expect(pageLoadTraceId).toEqual(serverTimingTraceId);
 });
 
-test('No sentry-trace meta tag is present (testing Server-Timing-only propagation)', async ({ page }) => {
+test('Test app does not render sentry-trace meta tags (precondition for Server-Timing tests)', async ({ page }) => {
   await page.goto('/');
 
   // Verify that NO sentry-trace meta tag is present
@@ -136,3 +139,115 @@ test('No sentry-trace meta tag is present (testing Server-Timing-only propagatio
   expect(sentryTraceMetaTag).toBeNull();
   expect(baggageMetaTag).toBeNull();
 });
+
+test('Propagates baggage/DSC from server to client via Server-Timing header', async ({ page }) => {
+  const testTag = crypto.randomUUID();
+
+  const responsePromise = page.waitForResponse(
+    response => response.url().includes(`tag=${testTag}`) && response.status() === 200,
+  );
+
+  const pageLoadTransactionPromise = waitForTransaction('remix-server-timing', transactionEvent => {
+    return transactionEvent.contexts?.trace?.op === 'pageload' && transactionEvent.tags?.['sentry_test'] === testTag;
+  });
+
+  await page.goto(`/?tag=${testTag}`);
+
+  const response = await responsePromise;
+  const serverTimingHeader = response.headers()['server-timing'];
+
+  // Extract sentry-trace from Server-Timing header
+  const sentryTraceMatch = serverTimingHeader?.match(/sentry-trace;desc="([^"]+)"/);
+  expect(sentryTraceMatch).toBeTruthy();
+  const sentryTraceValue = sentryTraceMatch?.[1];
+  const [headerTraceId] = sentryTraceValue?.split('-') || [];
+
+  // Extract baggage from Server-Timing header
+  const baggageMatch = serverTimingHeader?.match(/baggage;desc="([^"]+)"/);
+  expect(baggageMatch).toBeTruthy();
+
+  // Baggage is URL-encoded in Server-Timing header
+  const encodedBaggage = baggageMatch?.[1];
+  const decodedBaggage = decodeURIComponent(encodedBaggage || '');
+
+  // Parse baggage string into key-value pairs
+  const baggageEntries = decodedBaggage.split(',').reduce(
+    (acc, entry) => {
+      const [key, value] = entry.split('=');
+      if (key && value) {
+        acc[key] = value;
+      }
+      return acc;
+    },
+    {} as Record<string, string>,
+  );
+
+  // Verify essential DSC fields are present in baggage
+  expect(baggageEntries['sentry-environment']).toBeDefined();
+  expect(baggageEntries['sentry-trace_id']).toBeDefined();
+  expect(baggageEntries['sentry-public_key']).toBeDefined();
+
+  // CRITICAL: The trace_id in baggage MUST match the trace_id in sentry-trace header
+  // Both are generated from the same span in generateSentryServerTimingHeader()
+  // If they differ, there's a bug in DSC/trace context handling
+  expect(baggageEntries['sentry-trace_id']).toEqual(headerTraceId);
+
+  const pageloadTransaction = await pageLoadTransactionPromise;
+
+  // Verify the client transaction uses the trace_id from sentry-trace header
+  expect(pageloadTransaction.contexts?.trace?.trace_id).toEqual(headerTraceId);
+});
+
+test('Client pageload continues server trace with correct parent span ID', async ({ page }) => {
+  // This test verifies the complete trace chain:
+  // 1. Server generates Server-Timing header with trace_id and span_id
+  // 2. Client receives and parses the header
+  // 3. Client pageload transaction uses the same trace_id
+  // 4. Client pageload's parent_span_id matches the server's span_id
+  //
+  // Note: We verify trace continuity via the Server-Timing header rather than
+  // waiting for the server transaction, as the header is the propagation mechanism.
+
+  const testTag = crypto.randomUUID();
+
+  const responsePromise = page.waitForResponse(
+    response => response.url().includes(`tag=${testTag}`) && response.status() === 200,
+  );
+
+  const clientTransactionPromise = waitForTransaction('remix-server-timing', transactionEvent => {
+    return transactionEvent.contexts?.trace?.op === 'pageload' && transactionEvent.tags?.['sentry_test'] === testTag;
+  });
+
+  await page.goto(`/?tag=${testTag}`);
+
+  const response = await responsePromise;
+  const serverTimingHeader = response.headers()['server-timing'];
+
+  // Extract trace info from Server-Timing header (this represents the server span)
+  const sentryTraceMatch = serverTimingHeader?.match(/sentry-trace;desc="([^"]+)"/);
+  expect(sentryTraceMatch).toBeTruthy();
+  const serverTimingTraceValue = sentryTraceMatch?.[1];
+  const [serverTraceId, serverSpanId, sampled] = serverTimingTraceValue?.split('-') || [];
+
+  // Verify server trace info is valid
+  expect(serverTraceId).toBeDefined();
+  expect(serverTraceId).toHaveLength(32);
+  expect(serverSpanId).toBeDefined();
+  expect(serverSpanId).toHaveLength(16);
+  expect(sampled).toBe('1'); // Should be sampled
+
+  const clientTransaction = await clientTransactionPromise;
+
+  // Verify client transaction
+  expect(clientTransaction).toBeDefined();
+  const clientTraceId = clientTransaction.contexts?.trace?.trace_id;
+  const clientParentSpanId = clientTransaction.contexts?.trace?.parent_span_id;
+
+  // CRITICAL: Client trace_id must match server trace_id
+  // This proves trace propagation worked
+  expect(clientTraceId).toEqual(serverTraceId);
+
+  // CRITICAL: Client's parent_span_id must match the server's span_id
+  // This proves the client is continuing from the exact server span that sent the header
+  expect(clientParentSpanId).toEqual(serverSpanId);
+});

packages/remix/src/server/serverTimingTracePropagation.ts

@@ -1,7 +1,18 @@
 import type { Span } from '@sentry/core';
-import { debug, getActiveSpan, getRootSpan, getTraceData, isNodeEnv, spanToTraceHeader } from '@sentry/core';
+import {
+  debug,
+  getActiveSpan,
+  getDynamicSamplingContextFromSpan,
+  getRootSpan,
+  getTraceData,
+  isNodeEnv,
+  spanToTraceHeader,
+} from '@sentry/core';
 import { DEBUG_BUILD } from '../utils/debug-build';
 
+// Sentry baggage key prefix
+const SENTRY_BAGGAGE_KEY_PREFIX = 'sentry-';
+
 export interface ServerTimingTraceOptions {
   /** Include baggage in Server-Timing header. @default true */
   includeBaggage?: boolean;
@@ -46,8 +57,22 @@ export function generateSentryServerTimingHeader(options: ServerTimingTraceOptio
 
   if (span) {
     sentryTrace = spanToTraceHeader(span);
-    const traceData = getTraceData({ span });
-    baggage = traceData.baggage;
+    const spanTraceId = span.spanContext().traceId;
+
+    // Get DSC from span and ensure trace_id consistency
+    const dsc = getDynamicSamplingContextFromSpan(span);
+
+    // Build baggage string, ensuring trace_id matches the span's trace_id
+    // The DSC may have a different trace_id if it was frozen from an earlier context
+    const baggageEntries: string[] = [];
+    for (const [key, value] of Object.entries(dsc)) {
+      if (value) {
+        // Override trace_id to match the span's trace_id for consistency
+        const actualValue = key === 'trace_id' ? spanTraceId : value;
+        baggageEntries.push(`${SENTRY_BAGGAGE_KEY_PREFIX}${key}=${actualValue}`);
+      }
+    }
+    baggage = baggageEntries.length > 0 ? baggageEntries.join(',') : undefined;
   } else {
     const traceData = getTraceData();
     sentryTrace = traceData['sentry-trace'];