Add support for async context propagation and baggage in tracing service.

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

packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.test.ts

@@ -13,7 +13,13 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { SpanKind, SpanStatusCode, trace } from '@opentelemetry/api';
+import {
+  SpanKind,
+  SpanStatusCode,
+  context,
+  propagation,
+  trace,
+} from '@opentelemetry/api';
 import { mockCredentials, mockServices } from '@backstage/backend-test-utils';
 import { DefaultTracingService } from './DefaultTracingService';
 
@@ -267,4 +273,71 @@ describe('DefaultTracingService', () => {
     expect(value).toBe(42);
     expect(mocks.span.end).toHaveBeenCalledTimes(1);
   });
+
+  describe('withPropagatedContext', () => {
+    it('extracts context from headers and runs fn within it', async () => {
+      const extractSpy = jest.spyOn(propagation, 'extract');
+      const withSpy = jest.spyOn(context, 'with');
+
+      const service = createService();
+      const headers = { traceparent: '00-abc-def-01' };
+      const result = await service.withPropagatedContext(headers, () => 42);
+
+      expect(result).toBe(42);
+      expect(extractSpy).toHaveBeenCalledWith(expect.anything(), headers);
+      expect(withSpy).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.any(Function),
+      );
+    });
+
+    it('returns the value from an async fn', async () => {
+      const service = createService();
+      const result = await service.withPropagatedContext(
+        {},
+        async () => 'async-val',
+      );
+      expect(result).toBe('async-val');
+    });
+  });
+
+  describe('getActiveBaggage', () => {
+    it('returns undefined when no baggage is present', () => {
+      jest.spyOn(propagation, 'getActiveBaggage').mockReturnValue(undefined);
+      const service = createService();
+      expect(service.getActiveBaggage()).toBeUndefined();
+    });
+
+    it('returns a read-only baggage wrapping the active context baggage', () => {
+      const mockBaggage = {
+        getEntry: jest.fn((key: string) =>
+          key === 'gen_ai.conversation.id' ? { value: 'conv-1' } : undefined,
+        ),
+        getAllEntries: jest.fn(() => [
+          ['gen_ai.conversation.id', { value: 'conv-1' }],
+          ['gen_ai.agent.id', { value: 'agent-2' }],
+        ]),
+        setEntry: jest.fn(),
+        removeEntry: jest.fn(),
+        removeEntries: jest.fn(),
+        clear: jest.fn(),
+      };
+      jest
+        .spyOn(propagation, 'getActiveBaggage')
+        .mockReturnValue(mockBaggage as any);
+
+      const service = createService();
+      const baggage = service.getActiveBaggage();
+
+      expect(baggage).toBeDefined();
+      expect(baggage!.getEntry('gen_ai.conversation.id')).toEqual({
+        value: 'conv-1',
+      });
+      expect(baggage!.getEntry('unknown')).toBeUndefined();
+      expect(baggage!.getAllEntries()).toEqual([
+        ['gen_ai.conversation.id', { value: 'conv-1' }],
+        ['gen_ai.agent.id', { value: 'agent-2' }],
+      ]);
+    });
+  });
 });

packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts

@@ -14,7 +14,14 @@
  * limitations under the License.
  */
 
-import { SpanKind, SpanStatusCode, Tracer, trace } from '@opentelemetry/api';
+import {
+  SpanKind,
+  SpanStatusCode,
+  Tracer,
+  context,
+  propagation,
+  trace,
+} from '@opentelemetry/api';
 import {
   BackstageCredentials,
   HttpAuthService,
@@ -117,6 +124,29 @@ export class DefaultTracingService implements TracingService {
     );
   }
 
+  async withPropagatedContext<T>(
+    headers: Record<string, string | string[] | undefined>,
+    fn: () => T | Promise<T>,
+  ): Promise<T> {
+    const otelCtx = propagation.extract(context.active(), headers);
+    return context.with(otelCtx, fn);
+  }
+
+  getActiveBaggage() {
+    const baggage = propagation.getActiveBaggage();
+    if (!baggage) return undefined;
+    return {
+      getEntry: (key: string) => {
+        const entry = baggage.getEntry(key);
+        return entry ? { value: entry.value } : undefined;
+      },
+      getAllEntries: (): Array<[string, { value: string }]> =>
+        baggage
+          .getAllEntries()
+          .map(([key, entry]) => [key, { value: entry.value }]),
+    };
+  }
+
   private getPrincipalAttributes(
     credentials: BackstageCredentials | undefined,
   ): TracingServiceAttributes {

packages/backend-plugin-api/report-alpha.api.md

@@ -292,11 +292,16 @@ export const rootSystemMetadataServiceRef: ServiceRef<
 
 // @alpha
 export interface TracingService {
+  getActiveBaggage(): TracingServiceBaggage | undefined;
   startActiveSpan<T>(
     name: string,
     fn: (span: TracingServiceSpan) => T | Promise<T>,
     options?: TracingServiceSpanOptions,
   ): Promise<T>;
+  withPropagatedContext<T>(
+    headers: Record<string, string | string[] | undefined>,
+    fn: () => T | Promise<T>,
+  ): Promise<T>;
 }
 
 // @alpha
@@ -314,6 +319,20 @@ export type TracingServiceAttributeValue =
   | Array<null | undefined | number>
   | Array<null | undefined | boolean>;
 
+// @alpha
+export interface TracingServiceBaggage {
+  // (undocumented)
+  getAllEntries(): Array<[string, TracingServiceBaggageEntry]>;
+  // (undocumented)
+  getEntry(key: string): TracingServiceBaggageEntry | undefined;
+}
+
+// @alpha
+export interface TracingServiceBaggageEntry {
+  // (undocumented)
+  value: string;
+}
+
 // @alpha
 export const tracingServiceRef: ServiceRef<
   TracingService,

packages/backend-plugin-api/src/alpha/TracingService.ts

@@ -106,4 +106,39 @@ export interface TracingService {
     fn: (span: TracingServiceSpan) => T | Promise<T>,
     options?: TracingServiceSpanOptions,
   ): Promise<T>;
+
+  /**
+   * Extracts propagated context from HTTP headers and runs `fn` within
+   * it. Use this to bridge context across async boundaries where
+   * automatic propagation is lost.
+   */
+  withPropagatedContext<T>(
+    headers: Record<string, string | string[] | undefined>,
+    fn: () => T | Promise<T>,
+  ): Promise<T>;
+
+  /**
+   * Returns the active baggage from the current context, or `undefined`
+   * when none is present.
+   */
+  getActiveBaggage(): TracingServiceBaggage | undefined;
+}
+
+/**
+ * A read-only view of propagated baggage entries.
+ *
+ * @alpha
+ */
+export interface TracingServiceBaggage {
+  getEntry(key: string): TracingServiceBaggageEntry | undefined;
+  getAllEntries(): Array<[string, TracingServiceBaggageEntry]>;
+}
+
+/**
+ * A single baggage entry.
+ *
+ * @alpha
+ */
+export interface TracingServiceBaggageEntry {
+  value: string;
 }

packages/backend-plugin-api/src/alpha/index.ts

@@ -50,6 +50,8 @@ export type {
   TracingService,
   TracingServiceAttributeValue,
   TracingServiceAttributes,
+  TracingServiceBaggage,
+  TracingServiceBaggageEntry,
   TracingServiceSpan,
   TracingServiceSpanKind,
   TracingServiceSpanOptions,

packages/backend-plugin-api/src/alpha/refs.ts

@@ -71,7 +71,7 @@ export const metricsServiceRef = createServiceRef<
 /**
  * Service for managing trace spans.
  *
- * See {@link TracingService} for the API surface.
+ * See `TracingService` for the API surface.
  *
  * @alpha
  */

packages/backend-test-utils/report-alpha.api.md

@@ -108,9 +108,15 @@ export type ServiceMock<TService> = {
 export interface TracingServiceMock extends TracingService {
   // (undocumented)
   factory: ServiceFactory<TracingService>;
+  // (undocumented)
+  getActiveBaggage: jest.MockedFunction<TracingService['getActiveBaggage']>;
   spans: MockedTracingServiceSpan[];
   // (undocumented)
   startActiveSpan: jest.MockedFunction<TracingService['startActiveSpan']>;
+  // (undocumented)
+  withPropagatedContext: jest.MockedFunction<
+    TracingService['withPropagatedContext']
+  >;
 }
 
 // @alpha (undocumented)

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

@@ -46,6 +46,10 @@ export interface MockedTracingServiceSpan extends TracingServiceSpan {
  */
 export interface TracingServiceMock extends TracingService {
   startActiveSpan: jest.MockedFunction<TracingService['startActiveSpan']>;
+  withPropagatedContext: jest.MockedFunction<
+    TracingService['withPropagatedContext']
+  >;
+  getActiveBaggage: jest.MockedFunction<TracingService['getActiveBaggage']>;
   /** Spans created by `startActiveSpan` calls, in order. */
   spans: MockedTracingServiceSpan[];
   factory: ServiceFactory<TracingService>;
@@ -75,7 +79,18 @@ export namespace tracingServiceMock {
       return await fn(span);
     }) as TracingServiceMock['startActiveSpan'];
 
-    const service: TracingService = { startActiveSpan };
+    const withPropagatedContext = jest.fn(async (_headers, fn) =>
+      fn(),
+    ) as TracingServiceMock['withPropagatedContext'];
+    const getActiveBaggage = jest.fn(
+      () => undefined,
+    ) as TracingServiceMock['getActiveBaggage'];
+
+    const service: TracingService = {
+      startActiveSpan,
+      withPropagatedContext,
+      getActiveBaggage,
+    };
 
     return Object.assign(service as TracingServiceMock, {
       spans,