Simplify baggage getting to a single method
Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
+3
-9
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user