Simplify baggage getting to a single method

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
Eric Peterson
2026-05-18 16:50:43 +02:00
parent 3914351b84
commit 00bdd871a3
8 changed files with 23 additions and 64 deletions
@@ -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', () => {
@@ -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()
@@ -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
@@ -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]>;
}
@@ -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 }]
@@ -73,7 +73,6 @@ function parseBaggageHeader(
if (entries.size === 0) return undefined;
return {
getEntry: key => entries.get(key),
getAllEntries: () => Array.from(entries.entries()),
};
}
@@ -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<string, { value: string }> = {
'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<string, { value: string }> = {
'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' }],
@@ -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<string> = 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<string, string> = {};
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;
}
}