Propagate and forward gen_ai baggage from caller, if available

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>

packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2026 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { tracingServiceMock } from './TracingServiceMock';
+
+describe('tracingServiceMock', () => {
+  it('parses the baggage header from withPropagatedContext and exposes it via getActiveBaggage', async () => {
+    const tracing = tracingServiceMock.mock();
+
+    const seen = await tracing.withPropagatedContext(
+      {
+        baggage:
+          'gen_ai.conversation.id=conv-123, gen_ai.agent.id=agent-456;property=ignored',
+      },
+      () => {
+        const baggage = tracing.getActiveBaggage();
+        return {
+          conv: baggage?.getEntry('gen_ai.conversation.id')?.value,
+          agent: baggage?.getEntry('gen_ai.agent.id')?.value,
+          missing: baggage?.getEntry('gen_ai.missing'),
+          all: baggage?.getAllEntries().map(([k, v]) => [k, v.value]),
+        };
+      },
+    );
+
+    expect(seen).toEqual({
+      conv: 'conv-123',
+      agent: 'agent-456',
+      missing: undefined,
+      all: [
+        ['gen_ai.conversation.id', 'conv-123'],
+        ['gen_ai.agent.id', 'agent-456'],
+      ],
+    });
+
+    // Baggage is scoped to the propagated callback.
+    expect(tracing.getActiveBaggage()).toBeUndefined();
+  });
+
+  it('honours mockReturnValue overrides for getActiveBaggage', async () => {
+    const tracing = tracingServiceMock.mock();
+    const override = {
+      getEntry: () => ({ value: 'override' }),
+      getAllEntries: () =>
+        [['gen_ai.conversation.id', { value: 'override' }]] as Array<
+          [string, { value: string }]
+        >,
+    };
+    tracing.getActiveBaggage.mockReturnValue(override);
+
+    expect(tracing.getActiveBaggage()).toBe(override);
+    await tracing.withPropagatedContext(
+      { baggage: 'gen_ai.conversation.id=conv-from-header' },
+      () => {
+        // mockReturnValue takes precedence over the default header parsing.
+        expect(tracing.getActiveBaggage()).toBe(override);
+      },
+    );
+  });
+
+  it('returns undefined baggage when no baggage header is supplied', async () => {
+    const tracing = tracingServiceMock.mock();
+    await tracing.withPropagatedContext({ traceparent: 'whatever' }, () => {
+      expect(tracing.getActiveBaggage()).toBeUndefined();
+    });
+  });
+});

packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts

@@ -21,12 +21,47 @@ import {
 import {
   TracingService,
   TracingServiceAttributeValue,
+  TracingServiceBaggage,
   TracingServiceSpan,
   TracingServiceSpanStatus,
   tracingServiceRef,
 } from '@backstage/backend-plugin-api/alpha';
 import { tracingServiceFactory } from '@backstage/backend-defaults/alpha';
 
+// Parses the `baggage` header per the W3C Baggage member syntax,
+// dropping value properties (`;property=value`). This mirrors what
+// `propagation.extract` does in the real tracing service, just enough
+// for tests to assert end-to-end behaviour between propagated headers
+// and `getActiveBaggage()`.
+function parseBaggageHeader(
+  headers: Record<string, string | string[] | undefined>,
+): TracingServiceBaggage | undefined {
+  let raw: string | undefined;
+  for (const [name, value] of Object.entries(headers)) {
+    if (name.toLowerCase() !== 'baggage') continue;
+    raw = Array.isArray(value) ? value[0] : value;
+    break;
+  }
+  if (!raw) return undefined;
+
+  const entries = new Map<string, { value: string }>();
+  for (const segment of raw.split(',')) {
+    const [pair] = segment.split(';');
+    const eqIdx = pair.indexOf('=');
+    if (eqIdx === -1) continue;
+    const key = decodeURIComponent(pair.slice(0, eqIdx).trim());
+    const value = decodeURIComponent(pair.slice(eqIdx + 1).trim());
+    if (!key) continue;
+    entries.set(key, { value });
+  }
+  if (entries.size === 0) return undefined;
+
+  return {
+    getEntry: key => entries.get(key),
+    getAllEntries: () => Array.from(entries.entries()),
+  };
+}
+
 /**
  * A jest-mocked span captured by {@link TracingServiceMock}.
  *
@@ -42,6 +77,14 @@ export interface MockedTracingServiceSpan extends TracingServiceSpan {
  * `startActiveSpan` so tests can assert on the options passed in and the
  * methods called on the span inside the callback.
  *
+ * By default, `withPropagatedContext` parses the `baggage` header (W3C
+ * Baggage syntax) out of the supplied headers and makes those entries
+ * available via `getActiveBaggage` for the duration of the wrapped
+ * callback. Other propagation headers (e.g. `traceparent`) are ignored.
+ * Tests that need fully custom baggage can still override
+ * `getActiveBaggage` via `mockReturnValue` / `mockImplementation`, which
+ * takes precedence over the default behaviour.
+ *
  * @alpha
  */
 export interface TracingServiceMock extends TracingService {
@@ -79,11 +122,17 @@ export namespace tracingServiceMock {
       return await fn(span);
     }) as TracingServiceMock['startActiveSpan'];
 
-    const withPropagatedContext = jest.fn(async (_headers, fn) =>
-      fn(),
-    ) as TracingServiceMock['withPropagatedContext'];
+    const baggageStack: Array<TracingServiceBaggage | undefined> = [];
+    const withPropagatedContext = jest.fn(async (headers, fn) => {
+      baggageStack.push(parseBaggageHeader(headers));
+      try {
+        return await fn();
+      } finally {
+        baggageStack.pop();
+      }
+    }) as TracingServiceMock['withPropagatedContext'];
     const getActiveBaggage = jest.fn(
-      () => undefined,
+      () => baggageStack[baggageStack.length - 1],
     ) as TracingServiceMock['getActiveBaggage'];
 
     const service: TracingService = {