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 {