From 00bdd871a37dbc63c83f3ef840e11ff64a5252f0 Mon Sep 17 00:00:00 2001 From: Eric Peterson Date: Mon, 18 May 2026 16:50:43 +0200 Subject: [PATCH] Simplify baggage getting to a single method Signed-off-by: Eric Peterson --- .../tracing/DefaultTracingService.test.ts | 12 ++----- .../tracing/DefaultTracingService.ts | 4 --- .../backend-plugin-api/report-alpha.api.md | 4 +-- .../src/alpha/TracingService.ts | 5 +-- .../alpha/services/TracingServiceMock.test.ts | 36 ++++++++----------- .../src/alpha/services/TracingServiceMock.ts | 1 - .../src/services/McpService.test.ts | 16 --------- .../src/services/McpService.ts | 9 +++-- 8 files changed, 23 insertions(+), 64 deletions(-) diff --git a/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.test.ts b/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.test.ts index 61332a58b0..a03e7bc031 100644 --- a/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.test.ts +++ b/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.test.ts @@ -351,13 +351,11 @@ describe('DefaultTracingService', () => { describe('getActiveBaggage', () => { 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' }], ]), + getEntry: jest.fn(), setEntry: jest.fn(), removeEntry: jest.fn(), removeEntries: jest.fn(), @@ -371,10 +369,6 @@ describe('DefaultTracingService', () => { const baggage = service.propagation.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' }], @@ -386,8 +380,8 @@ describe('DefaultTracingService', () => { it('returns baggage from the supplied context', () => { const ctx = { __ctx: 'has-baggage' } as any; const mockBaggage = { - getEntry: jest.fn(() => ({ value: 'ctx-val' })), getAllEntries: jest.fn(() => [['k', { value: 'ctx-val' }]]), + getEntry: jest.fn(), setEntry: jest.fn(), removeEntry: jest.fn(), removeEntries: jest.fn(), @@ -401,7 +395,7 @@ describe('DefaultTracingService', () => { const baggage = service.propagation.getBaggage(ctx); expect(getBaggageSpy).toHaveBeenCalledWith(ctx); - expect(baggage!.getEntry('k')).toEqual({ value: 'ctx-val' }); + expect(baggage!.getAllEntries()).toEqual([['k', { value: 'ctx-val' }]]); }); it('returns undefined when the context has no baggage', () => { diff --git a/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts b/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts index 48e30fdc2c..c58306f8ae 100644 --- a/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts +++ b/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts @@ -69,10 +69,6 @@ function wrapOtelBaggage( ): TracingServiceBaggage | undefined { if (!baggage) return undefined; return { - getEntry: (key: string) => { - const entry = baggage.getEntry(key); - return entry ? { value: entry.value } : undefined; - }, getAllEntries: () => baggage .getAllEntries() diff --git a/packages/backend-plugin-api/report-alpha.api.md b/packages/backend-plugin-api/report-alpha.api.md index 41a9d6160d..53ebf99f51 100644 --- a/packages/backend-plugin-api/report-alpha.api.md +++ b/packages/backend-plugin-api/report-alpha.api.md @@ -324,8 +324,6 @@ export type TracingServiceAttributeValue = export interface TracingServiceBaggage { // (undocumented) getAllEntries(): Array<[string, TracingServiceBaggageEntry]>; - // (undocumented) - getEntry(key: string): TracingServiceBaggageEntry | undefined; } // @alpha @@ -337,7 +335,7 @@ export interface TracingServiceBaggageEntry { // @alpha export interface TracingServiceContext { // (undocumented) - readonly [tracingServiceContextBrand]: never; + readonly $$type: '@backstage/TracingServiceContext'; } // @alpha diff --git a/packages/backend-plugin-api/src/alpha/TracingService.ts b/packages/backend-plugin-api/src/alpha/TracingService.ts index f751b776bd..3b159f975c 100644 --- a/packages/backend-plugin-api/src/alpha/TracingService.ts +++ b/packages/backend-plugin-api/src/alpha/TracingService.ts @@ -176,8 +176,6 @@ export interface TracingServicePropagationAPI { getActiveBaggage(): TracingServiceBaggage | undefined; } -declare const tracingServiceContextBrand: unique symbol; - /** * Opaque handle representing a tracing context. Consumers receive * these from {@link TracingServiceContextAPI.active} or @@ -187,7 +185,7 @@ declare const tracingServiceContextBrand: unique symbol; * @alpha */ export interface TracingServiceContext { - readonly [tracingServiceContextBrand]: never; + readonly $$type: '@backstage/TracingServiceContext'; } /** @@ -196,7 +194,6 @@ export interface TracingServiceContext { * @alpha */ export interface TracingServiceBaggage { - getEntry(key: string): TracingServiceBaggageEntry | undefined; getAllEntries(): Array<[string, TracingServiceBaggageEntry]>; } diff --git a/packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts b/packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts index 5fb711a636..7e3cd7fa92 100644 --- a/packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts +++ b/packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts @@ -26,29 +26,22 @@ describe('tracingServiceMock', () => { }); // Baggage is reachable directly off the extracted handle. - expect( - tracing.propagation.getBaggage(ctx)?.getEntry('gen_ai.agent.id'), - ).toEqual({ value: 'agent-456' }); + expect(tracing.propagation.getBaggage(ctx)?.getAllEntries()).toEqual([ + ['gen_ai.conversation.id', { value: 'conv-123' }], + ['gen_ai.agent.id', { value: 'agent-456' }], + ]); - const seen = await tracing.context.with(ctx, () => { - const baggage = tracing.propagation.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]), - }; - }); + const seen = await tracing.context.with(ctx, () => + tracing.propagation + .getActiveBaggage() + ?.getAllEntries() + .map(([k, v]) => [k, v.value]), + ); - expect(seen).toEqual({ - conv: 'conv-123', - agent: 'agent-456', - missing: undefined, - all: [ - ['gen_ai.conversation.id', { value: 'conv-123' }], - ['gen_ai.agent.id', { value: 'agent-456' }], - ].map(([k, v]) => [k, (v as { value: string }).value]), - }); + expect(seen).toEqual([ + ['gen_ai.conversation.id', 'conv-123'], + ['gen_ai.agent.id', 'agent-456'], + ]); // Baggage is scoped to the context.with callback. expect(tracing.propagation.getActiveBaggage()).toBeUndefined(); @@ -57,7 +50,6 @@ describe('tracingServiceMock', () => { 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 }] diff --git a/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts b/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts index 69125a7a2f..b964e09b33 100644 --- a/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts +++ b/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts @@ -73,7 +73,6 @@ function parseBaggageHeader( if (entries.size === 0) return undefined; return { - getEntry: key => entries.get(key), getAllEntries: () => Array.from(entries.entries()), }; } diff --git a/plugins/mcp-actions-backend/src/services/McpService.test.ts b/plugins/mcp-actions-backend/src/services/McpService.test.ts index 7278bf276e..8a84bd9198 100644 --- a/plugins/mcp-actions-backend/src/services/McpService.test.ts +++ b/plugins/mcp-actions-backend/src/services/McpService.test.ts @@ -1013,13 +1013,6 @@ describe('McpService', () => { it('includes gen_ai baggage entries as span attributes when present', async () => { const tracing = tracingServiceMock.mock(); tracing.propagation.getActiveBaggage.mockReturnValue({ - getEntry: (key: string) => { - const entries: Record = { - 'gen_ai.conversation.id': { value: 'conv-123' }, - 'gen_ai.agent.id': { value: 'agent-456' }, - }; - return entries[key]; - }, getAllEntries: () => [ ['gen_ai.conversation.id', { value: 'conv-123' }], ['gen_ai.agent.id', { value: 'agent-456' }], @@ -1036,15 +1029,6 @@ describe('McpService', () => { it('only forwards allowlisted baggage keys onto the span', async () => { const tracing = tracingServiceMock.mock(); tracing.propagation.getActiveBaggage.mockReturnValue({ - getEntry: (key: string) => { - const entries: Record = { - 'gen_ai.conversation.id': { value: 'conv-123' }, - 'gen_ai.tool.call.result': { value: 'injected-result' }, - 'gen_ai.prompt': { value: 'injected-prompt' }, - 'gen_ai.user.message': { value: 'injected-user-message' }, - }; - return entries[key]; - }, getAllEntries: () => [ ['gen_ai.conversation.id', { value: 'conv-123' }], ['gen_ai.tool.call.result', { value: 'injected-result' }], diff --git a/plugins/mcp-actions-backend/src/services/McpService.ts b/plugins/mcp-actions-backend/src/services/McpService.ts index 09685fed1d..25e1a7002d 100644 --- a/plugins/mcp-actions-backend/src/services/McpService.ts +++ b/plugins/mcp-actions-backend/src/services/McpService.ts @@ -46,13 +46,13 @@ function safeStringify(value: unknown): string { // Baggage is propagated from untrusted callers, so we forward only an // explicit allowlist of low-cardinality identifier keys from the OTel // `gen_ai.*` registry. -const PROPAGATED_BAGGAGE_ATTRIBUTES: readonly string[] = [ +const PROPAGATED_BAGGAGE_ATTRIBUTES: ReadonlySet = new Set([ 'gen_ai.agent.id', 'gen_ai.agent.name', 'gen_ai.conversation.id', 'gen_ai.provider.name', 'gen_ai.request.model', -]; +]); function baggageAttributes( tracingService: TracingService, @@ -60,9 +60,8 @@ function baggageAttributes( const baggage = tracingService.propagation.getActiveBaggage(); if (!baggage) return {}; const attrs: Record = {}; - for (const key of PROPAGATED_BAGGAGE_ATTRIBUTES) { - const entry = baggage.getEntry(key); - if (entry) { + for (const [key, entry] of baggage.getAllEntries()) { + if (PROPAGATED_BAGGAGE_ATTRIBUTES.has(key)) { attrs[key] = entry.value; } }