diff --git a/.changeset/tracing-service-context-and-baggage.md b/.changeset/tracing-service-context-and-baggage.md index 3f447cc339..86d1a62168 100644 --- a/.changeset/tracing-service-context-and-baggage.md +++ b/.changeset/tracing-service-context-and-baggage.md @@ -4,4 +4,4 @@ '@backstage/backend-test-utils': patch --- -Added `withPropagatedContext` and `getActiveBaggage` to the alpha `TracingService`, enabling plugins to bridge OpenTelemetry context across async boundaries and read propagated baggage. +Added `context` and `propagation` to the alpha `TracingService`. Plugins can bridge OpenTelemetry context across async boundaries via `tracing.propagation.extract(tracing.context.active(), carrier)` followed by `tracing.context.with(ctx, fn)`, and read propagated baggage via `tracing.propagation.getActiveBaggage()` or `tracing.propagation.getBaggage(ctx)`. diff --git a/docs/backend-system/core-services/tracing.md b/docs/backend-system/core-services/tracing.md index 3cd9b7c00f..52d2d3069f 100644 --- a/docs/backend-system/core-services/tracing.md +++ b/docs/backend-system/core-services/tracing.md @@ -102,30 +102,45 @@ The span object exposes: ## Context Propagation -When your plugin receives requests through a protocol layer that breaks automatic OpenTelemetry context propagation (e.g. a WebSocket handler), use `withPropagatedContext` to extract the trace parent and baggage from the incoming HTTP headers and run the handler within that context: +The tracing service exposes two sub-objects that mirror the corresponding namespaces in `@opentelemetry/api`: + +- `tracing.context` for context management (`active`, `with`). +- `tracing.propagation` for context propagation (`extract`, `getBaggage`, `getActiveBaggage`). + +When your plugin handles a request through a transport or framework that doesn't automatically attach the caller's context to the work it runs (for example, a handler dispatched from a message-queue consumer, or a third-party transport like the MCP streamable HTTP transport that re-enters user code outside of Express's middleware chain), extract the trace parent and baggage from the inbound request's headers yourself and run the handler with that context active: ```ts router.post('/', async (req, res) => { - await tracing.withPropagatedContext(req.headers, () => + const ctx = tracing.propagation.extract( + tracing.context.active(), + req.headers, + ); + await tracing.context.with(ctx, () => transport.handleRequest(req, res, req.body), ); }); ``` +`propagation.extract` reads from a header-shaped record (`Record`), so any source of headers — Express's `req.headers`, a Node.js `http.IncomingMessage`, or a payload field carrying serialized headers — works the same way. + Any spans created inside the callback — including those from `startActiveSpan` — will be children of the propagated trace and will have access to the propagated baggage. +The context returned by `propagation.extract` and `context.active` is an opaque handle: consumers pass it back into the API but do not introspect it. + ## Reading Baggage -Use `getActiveBaggage()` to read baggage entries from the active context. This is useful for forwarding caller-set metadata onto your spans — for example, a request ID, tenant identifier, or feature-flag context that the caller propagated via baggage: +Use `propagation.getActiveBaggage()` to read baggage entries from the currently active context. This is useful for forwarding caller-set metadata onto your spans — for example, a request ID, tenant identifier, or feature-flag context that the caller propagated via baggage: ```ts -const baggage = tracing.getActiveBaggage(); +const baggage = tracing.propagation.getActiveBaggage(); const tenantId = baggage?.getEntry('app.tenant.id')?.value; if (tenantId) { span.setAttribute('app.tenant.id', tenantId); } ``` +Use `propagation.getBaggage(ctx)` when you already hold a specific context handle (for example, one returned by `propagation.extract`) and want to read its baggage without first activating the context. + The returned object exposes: | Method | Description | @@ -133,7 +148,7 @@ The returned object exposes: | `getEntry(key)` | Returns `{ value: string }` for the key, or `undefined`. | | `getAllEntries()` | Returns all entries as `[key, { value }][]`. | -Returns `undefined` when no baggage is present in the active context. +Both calls return `undefined` when no baggage is present. ## Principal Enrichment 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 c7dcbd9925..61332a58b0 100644 --- a/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.test.ts +++ b/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.test.ts @@ -94,10 +94,11 @@ describe('DefaultTracingService', () => { it('passes name, kind, and caller attributes through to the tracer', async () => { const service = createService(); - await service.startActiveSpan('op', async () => undefined, { - kind: 'server', - attributes: { foo: 'bar' }, - }); + await service.startActiveSpan( + 'op', + { kind: 'server', attributes: { foo: 'bar' } }, + async () => undefined, + ); expect(mocks.tracer.startActiveSpan).toHaveBeenCalledWith( 'op', @@ -119,9 +120,11 @@ describe('DefaultTracingService', () => { it('lets caller-supplied attributes override backstage.plugin.id at start time', async () => { const service = createService({ pluginId: 'my-plugin' }); - await service.startActiveSpan('op', async () => undefined, { - attributes: { 'backstage.plugin.id': 'other-plugin' }, - }); + await service.startActiveSpan( + 'op', + { attributes: { 'backstage.plugin.id': 'other-plugin' } }, + async () => undefined, + ); const attrs = mocks.tracer.startActiveSpan.mock.calls[0][1].attributes; expect(attrs['backstage.plugin.id']).toBe('other-plugin'); @@ -146,18 +149,22 @@ describe('DefaultTracingService', () => { ]; for (const [kind, expected] of cases) { mocks.tracer.startActiveSpan.mockClear(); - await service.startActiveSpan('op', async () => undefined, { - kind: kind as any, - }); + await service.startActiveSpan( + 'op', + { kind: kind as any }, + async () => undefined, + ); expect(mocks.tracer.startActiveSpan.mock.calls[0][1].kind).toBe(expected); } }); it('adds backstage.principal.type but not enduser.id when capture is off', async () => { const service = createService({ captureEndUser: false }); - await service.startActiveSpan('op', async () => undefined, { - credentials: mockCredentials.user('user:default/alice'), - }); + await service.startActiveSpan( + 'op', + { credentials: mockCredentials.user('user:default/alice') }, + async () => undefined, + ); const attrs = mocks.tracer.startActiveSpan.mock.calls[0][1].attributes; expect(attrs['backstage.principal.type']).toBe('user'); @@ -166,9 +173,11 @@ describe('DefaultTracingService', () => { it('adds enduser.id from a user principal when capture is on', async () => { const service = createService({ captureEndUser: true }); - await service.startActiveSpan('op', async () => undefined, { - credentials: mockCredentials.user('user:default/alice'), - }); + await service.startActiveSpan( + 'op', + { credentials: mockCredentials.user('user:default/alice') }, + async () => undefined, + ); const attrs = mocks.tracer.startActiveSpan.mock.calls[0][1].attributes; expect(attrs['enduser.id']).toBe('user:default/alice'); @@ -176,9 +185,11 @@ describe('DefaultTracingService', () => { it('adds enduser.id from a service principal subject when capture is on', async () => { const service = createService({ captureEndUser: true }); - await service.startActiveSpan('op', async () => undefined, { - credentials: mockCredentials.service('plugin:test'), - }); + await service.startActiveSpan( + 'op', + { credentials: mockCredentials.service('plugin:test') }, + async () => undefined, + ); const attrs = mocks.tracer.startActiveSpan.mock.calls[0][1].attributes; expect(attrs['enduser.id']).toBe('plugin:test'); @@ -195,9 +206,11 @@ describe('DefaultTracingService', () => { httpAuth, }); - await service.startActiveSpan('op', async () => undefined, { - request: { headers: {} } as any, - }); + await service.startActiveSpan( + 'op', + { request: { headers: {} } as any }, + async () => undefined, + ); expect(credSpy).toHaveBeenCalledTimes(1); const attrs = mocks.tracer.startActiveSpan.mock.calls[0][1].attributes; @@ -214,10 +227,14 @@ describe('DefaultTracingService', () => { httpAuth, }); - await service.startActiveSpan('op', async () => undefined, { - credentials: mockCredentials.user('user:default/explicit'), - request: { headers: {} } as any, - }); + await service.startActiveSpan( + 'op', + { + credentials: mockCredentials.user('user:default/explicit'), + request: { headers: {} } as any, + }, + async () => undefined, + ); expect(credSpy).not.toHaveBeenCalled(); const attrs = mocks.tracer.startActiveSpan.mock.calls[0][1].attributes; @@ -274,70 +291,124 @@ describe('DefaultTracingService', () => { 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), - ); + describe('context', () => { + describe('active', () => { + it('returns the OTel active context as an opaque handle', () => { + const fakeCtx = { __ctx: 'active' }; + jest.spyOn(context, 'active').mockReturnValue(fakeCtx as any); + const service = createService(); + expect(service.context.active()).toBe(fakeCtx); + }); }); - 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('with', () => { + it('delegates to OTel context.with on the supplied handle and returns the fn result', async () => { + const fakeCtx = { __ctx: 'extracted' } as any; + const withSpy = jest + .spyOn(context, 'with') + .mockImplementation((_ctx, fn) => (fn as any)()); + + const service = createService(); + const result = await service.context.with(fakeCtx, () => 99); + + expect(withSpy).toHaveBeenCalledWith(fakeCtx, expect.any(Function)); + expect(result).toBe(99); + }); + + it('awaits an async fn and returns its resolved value', async () => { + jest + .spyOn(context, 'with') + .mockImplementation((_ctx, fn) => (fn as any)()); + + const service = createService(); + const result = await service.context.with( + {} as any, + 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(); + describe('propagation', () => { + describe('extract', () => { + it('forwards the supplied context and headers to OTel propagation.extract', () => { + const baseCtx = { __ctx: 'base' } as any; + const extractedCtx = { __ctx: 'extracted' } as any; + const extractSpy = jest + .spyOn(propagation, 'extract') + .mockReturnValue(extractedCtx); + + const service = createService(); + const headers = { traceparent: '00-abc-def-01' }; + const result = service.propagation.extract(baseCtx, headers); + + expect(extractSpy).toHaveBeenCalledWith(baseCtx, headers); + expect(result).toBe(extractedCtx); + }); }); - 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(() => [ + 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' }], + ]), + 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.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' }], - ]), - 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', + ]); + }); + }); + + describe('getBaggage', () => { + 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' }]]), + setEntry: jest.fn(), + removeEntry: jest.fn(), + removeEntries: jest.fn(), + clear: jest.fn(), + }; + const getBaggageSpy = jest + .spyOn(propagation, 'getBaggage') + .mockReturnValue(mockBaggage as any); + + const service = createService(); + const baggage = service.propagation.getBaggage(ctx); + + expect(getBaggageSpy).toHaveBeenCalledWith(ctx); + expect(baggage!.getEntry('k')).toEqual({ value: 'ctx-val' }); + }); + + it('returns undefined when the context has no baggage', () => { + jest.spyOn(propagation, 'getBaggage').mockReturnValue(undefined); + const service = createService(); + expect(service.propagation.getBaggage({} as any)).toBeUndefined(); }); - expect(baggage!.getEntry('unknown')).toBeUndefined(); - expect(baggage!.getAllEntries()).toEqual([ - ['gen_ai.conversation.id', { value: 'conv-1' }], - ['gen_ai.agent.id', { value: 'agent-2' }], - ]); }); }); }); diff --git a/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts b/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts index 5d99562949..48e30fdc2c 100644 --- a/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts +++ b/packages/backend-defaults/src/alpha/entrypoints/tracing/DefaultTracingService.ts @@ -15,11 +15,12 @@ */ import { + Context, SpanKind, SpanStatusCode, Tracer, - context, - propagation, + context as otelContext, + propagation as otelPropagation, trace, } from '@opentelemetry/api'; import { @@ -29,6 +30,10 @@ import { import { TracingService, TracingServiceAttributes, + TracingServiceBaggage, + TracingServiceContext, + TracingServiceContextAPI, + TracingServicePropagationAPI, TracingServiceSpan, TracingServiceSpanKind, TracingServiceSpanOptions, @@ -49,6 +54,32 @@ export interface DefaultTracingServiceOptions { httpAuth: HttpAuthService; } +// `TracingServiceContext` is an opaque handle for an OTel `Context`. Internally +// the value *is* the OTel context; we just narrow the type so consumers can't +// poke at it directly. +function toOtelContext(ctx: TracingServiceContext): Context { + return ctx as unknown as Context; +} +function fromOtelContext(ctx: Context): TracingServiceContext { + return ctx as unknown as TracingServiceContext; +} + +function wrapOtelBaggage( + baggage: ReturnType, +): TracingServiceBaggage | undefined { + if (!baggage) return undefined; + return { + getEntry: (key: string) => { + const entry = baggage.getEntry(key); + return entry ? { value: entry.value } : undefined; + }, + getAllEntries: () => + baggage + .getAllEntries() + .map(([key, entry]) => [key, { value: entry.value }]), + }; +} + /** * Default implementation of the {@link TracingService} interface. * @@ -60,6 +91,25 @@ export class DefaultTracingService implements TracingService { private readonly captureEndUser: boolean; private readonly httpAuth: HttpAuthService; + readonly context: TracingServiceContextAPI = { + active: () => fromOtelContext(otelContext.active()), + with: async ( + ctx: TracingServiceContext, + fn: () => T | Promise, + ): Promise => otelContext.with(toOtelContext(ctx), fn), + }; + + readonly propagation: TracingServicePropagationAPI = { + extract: ( + ctx: TracingServiceContext, + carrier: Record, + ): TracingServiceContext => + fromOtelContext(otelPropagation.extract(toOtelContext(ctx), carrier)), + getBaggage: (ctx: TracingServiceContext) => + wrapOtelBaggage(otelPropagation.getBaggage(toOtelContext(ctx))), + getActiveBaggage: () => wrapOtelBaggage(otelPropagation.getActiveBaggage()), + }; + private constructor(opts: DefaultTracingServiceOptions) { this.tracer = trace .getTracerProvider() @@ -73,11 +123,30 @@ export class DefaultTracingService implements TracingService { return new DefaultTracingService(opts); } - async startActiveSpan( + startActiveSpan( name: string, fn: (span: TracingServiceSpan) => T | Promise, - options: TracingServiceSpanOptions = {}, + ): Promise; + startActiveSpan( + name: string, + options: TracingServiceSpanOptions, + fn: (span: TracingServiceSpan) => T | Promise, + ): Promise; + async startActiveSpan( + name: string, + optionsOrFn: + | TracingServiceSpanOptions + | ((span: TracingServiceSpan) => T | Promise), + maybeFn?: (span: TracingServiceSpan) => T | Promise, ): Promise { + const [options, fn]: [ + TracingServiceSpanOptions, + (span: TracingServiceSpan) => T | Promise, + ] = + typeof optionsOrFn === 'function' + ? [{}, optionsOrFn] + : [optionsOrFn, maybeFn!]; + let credentials = options.credentials; if (!credentials && options.request) { credentials = await this.httpAuth.credentials(options.request); @@ -124,29 +193,6 @@ export class DefaultTracingService implements TracingService { ); } - async withPropagatedContext( - headers: Record, - fn: () => T | Promise, - ): Promise { - 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 { diff --git a/packages/backend-plugin-api/report-alpha.api.md b/packages/backend-plugin-api/report-alpha.api.md index e970d85624..41a9d6160d 100644 --- a/packages/backend-plugin-api/report-alpha.api.md +++ b/packages/backend-plugin-api/report-alpha.api.md @@ -292,15 +292,16 @@ export const rootSystemMetadataServiceRef: ServiceRef< // @alpha export interface TracingService { - getActiveBaggage(): TracingServiceBaggage | undefined; + readonly context: TracingServiceContextAPI; + readonly propagation: TracingServicePropagationAPI; startActiveSpan( name: string, fn: (span: TracingServiceSpan) => T | Promise, - options?: TracingServiceSpanOptions, ): Promise; - withPropagatedContext( - headers: Record, - fn: () => T | Promise, + startActiveSpan( + name: string, + options: TracingServiceSpanOptions, + fn: (span: TracingServiceSpan) => T | Promise, ): Promise; } @@ -333,6 +334,28 @@ export interface TracingServiceBaggageEntry { value: string; } +// @alpha +export interface TracingServiceContext { + // (undocumented) + readonly [tracingServiceContextBrand]: never; +} + +// @alpha +export interface TracingServiceContextAPI { + active(): TracingServiceContext; + with(context: TracingServiceContext, fn: () => T | Promise): Promise; +} + +// @alpha +export interface TracingServicePropagationAPI { + extract( + context: TracingServiceContext, + carrier: Record, + ): TracingServiceContext; + getActiveBaggage(): TracingServiceBaggage | undefined; + getBaggage(context: TracingServiceContext): TracingServiceBaggage | undefined; +} + // @alpha export const tracingServiceRef: ServiceRef< TracingService, diff --git a/packages/backend-plugin-api/src/alpha/TracingService.ts b/packages/backend-plugin-api/src/alpha/TracingService.ts index 8919263676..f751b776bd 100644 --- a/packages/backend-plugin-api/src/alpha/TracingService.ts +++ b/packages/backend-plugin-api/src/alpha/TracingService.ts @@ -104,26 +104,92 @@ export interface TracingService { startActiveSpan( name: string, fn: (span: TracingServiceSpan) => T | Promise, - options?: TracingServiceSpanOptions, ): Promise; - /** - * Extracts propagated context from HTTP headers and runs `fn` within - * it. Use this to bridge context across async boundaries where - * automatic propagation is lost. + * Runs `fn` inside a new active span configured by `options`. The + * span is finished when `fn` resolves or throws. */ - withPropagatedContext( - headers: Record, - fn: () => T | Promise, + startActiveSpan( + name: string, + options: TracingServiceSpanOptions, + fn: (span: TracingServiceSpan) => T | Promise, ): Promise; /** - * Returns the active baggage from the current context, or `undefined` - * when none is present. + * Read the active tracing context, or run work within a specific + * one. + */ + readonly context: TracingServiceContextAPI; + + /** + * Extract a caller's tracing context from an inbound carrier, and + * read baggage from a context. Use these to bridge context across + * boundaries where automatic propagation is lost — for example, + * when a request arrives over a transport that does not + * automatically attach the caller's context. + */ + readonly propagation: TracingServicePropagationAPI; +} + +/** + * Read the active tracing context, or run work within a specific one. + * The context carries the active span and propagation fields (trace + * parent, baggage) for the current unit of work, and is automatically + * inherited by spans created via `startActiveSpan`. + * + * @alpha + */ +export interface TracingServiceContextAPI { + /** Returns the currently active context. */ + active(): TracingServiceContext; + /** Runs `fn` with the supplied context set as the active context. */ + with(context: TracingServiceContext, fn: () => T | Promise): Promise; +} + +/** + * Extract a caller's trace parent and baggage from an inbound carrier + * (typically HTTP headers), or read baggage from a context. Use these + * to bridge context across boundaries where automatic propagation is + * lost. + * + * @alpha + */ +export interface TracingServicePropagationAPI { + /** + * Returns a new context with propagation fields (trace parent, + * baggage, ...) read from the supplied carrier merged into it. + */ + extract( + context: TracingServiceContext, + carrier: Record, + ): TracingServiceContext; + /** + * Returns the baggage attached to the supplied context, or + * `undefined` when none is present. + */ + getBaggage(context: TracingServiceContext): TracingServiceBaggage | undefined; + /** + * Returns the baggage attached to the currently active context, or + * `undefined` when none is present. Equivalent to + * `getBaggage(context.active())`. */ getActiveBaggage(): TracingServiceBaggage | undefined; } +declare const tracingServiceContextBrand: unique symbol; + +/** + * Opaque handle representing a tracing context. Consumers receive + * these from {@link TracingServiceContextAPI.active} or + * {@link TracingServicePropagationAPI.extract} and pass them back into + * the API; the type carries no inspectable fields. + * + * @alpha + */ +export interface TracingServiceContext { + readonly [tracingServiceContextBrand]: never; +} + /** * A read-only view of propagated baggage entries. * diff --git a/packages/backend-plugin-api/src/alpha/index.ts b/packages/backend-plugin-api/src/alpha/index.ts index dc0624273b..1e2a03a04f 100644 --- a/packages/backend-plugin-api/src/alpha/index.ts +++ b/packages/backend-plugin-api/src/alpha/index.ts @@ -52,6 +52,9 @@ export type { TracingServiceAttributes, TracingServiceBaggage, TracingServiceBaggageEntry, + TracingServiceContext, + TracingServiceContextAPI, + TracingServicePropagationAPI, TracingServiceSpan, TracingServiceSpanKind, TracingServiceSpanOptions, diff --git a/packages/backend-test-utils/report-alpha.api.md b/packages/backend-test-utils/report-alpha.api.md index b2d6b56e81..f3faa4eb5e 100644 --- a/packages/backend-test-utils/report-alpha.api.md +++ b/packages/backend-test-utils/report-alpha.api.md @@ -16,6 +16,8 @@ import { MetricsService } from '@backstage/backend-plugin-api/alpha'; import { ServiceFactory } from '@backstage/backend-plugin-api'; import { TracingService } from '@backstage/backend-plugin-api/alpha'; import { TracingServiceAttributeValue } from '@backstage/backend-plugin-api/alpha'; +import { TracingServiceContextAPI } from '@backstage/backend-plugin-api/alpha'; +import { TracingServicePropagationAPI } from '@backstage/backend-plugin-api/alpha'; import { TracingServiceSpan } from '@backstage/backend-plugin-api/alpha'; import { TracingServiceSpanStatus } from '@backstage/backend-plugin-api/alpha'; @@ -85,6 +87,28 @@ export class MockActionsRegistry >(options: ActionsRegistryActionOptions): void; } +// @alpha +export interface MockedTracingServiceContextAPI + extends TracingServiceContextAPI { + // (undocumented) + active: jest.MockedFunction; + // (undocumented) + with: jest.MockedFunction; +} + +// @alpha +export interface MockedTracingServicePropagationAPI + extends TracingServicePropagationAPI { + // (undocumented) + extract: jest.MockedFunction; + // (undocumented) + getActiveBaggage: jest.MockedFunction< + TracingServicePropagationAPI['getActiveBaggage'] + >; + // (undocumented) + getBaggage: jest.MockedFunction; +} + // @alpha export interface MockedTracingServiceSpan extends TracingServiceSpan { // (undocumented) @@ -106,17 +130,15 @@ export type ServiceMock = { // @alpha export interface TracingServiceMock extends TracingService { + // (undocumented) + context: MockedTracingServiceContextAPI; // (undocumented) factory: ServiceFactory; // (undocumented) - getActiveBaggage: jest.MockedFunction; + propagation: MockedTracingServicePropagationAPI; spans: MockedTracingServiceSpan[]; // (undocumented) startActiveSpan: jest.MockedFunction; - // (undocumented) - withPropagatedContext: jest.MockedFunction< - TracingService['withPropagatedContext'] - >; } // @alpha (undocumented) 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 be80625932..5fb711a636 100644 --- a/packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts +++ b/packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts @@ -17,37 +17,41 @@ import { tracingServiceMock } from './TracingServiceMock'; describe('tracingServiceMock', () => { - it('parses the baggage header from withPropagatedContext and exposes it via getActiveBaggage', async () => { + it('parses the baggage header via propagation.extract and exposes it via getActiveBaggage inside context.with', async () => { const tracing = tracingServiceMock.mock(); - const seen = await tracing.withPropagatedContext( - { - baggage: - 'gen_ai.conversation.id=conv-123, gen_ai.agent.id=agent-456;property=ignored', - }, - () => { - const baggage = tracing.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 ctx = tracing.propagation.extract(tracing.context.active(), { + baggage: + 'gen_ai.conversation.id=conv-123, gen_ai.agent.id=agent-456;property=ignored', + }); + + // Baggage is reachable directly off the extracted handle. + expect( + tracing.propagation.getBaggage(ctx)?.getEntry('gen_ai.agent.id'), + ).toEqual({ 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]), + }; + }); expect(seen).toEqual({ conv: 'conv-123', agent: 'agent-456', missing: undefined, all: [ - ['gen_ai.conversation.id', 'conv-123'], - ['gen_ai.agent.id', 'agent-456'], - ], + ['gen_ai.conversation.id', { value: 'conv-123' }], + ['gen_ai.agent.id', { value: 'agent-456' }], + ].map(([k, v]) => [k, (v as { value: string }).value]), }); - // Baggage is scoped to the propagated callback. - expect(tracing.getActiveBaggage()).toBeUndefined(); + // Baggage is scoped to the context.with callback. + expect(tracing.propagation.getActiveBaggage()).toBeUndefined(); }); it('honours mockReturnValue overrides for getActiveBaggage', async () => { @@ -59,22 +63,25 @@ describe('tracingServiceMock', () => { [string, { value: string }] >, }; - tracing.getActiveBaggage.mockReturnValue(override); + tracing.propagation.getActiveBaggage.mockReturnValue(override); - expect(tracing.getActiveBaggage()).toBe(override); - await tracing.withPropagatedContext( - { baggage: 'gen_ai.conversation.id=conv-from-header' }, - () => { - // mockReturnValue takes precedence over the default header parsing. - expect(tracing.getActiveBaggage()).toBe(override); - }, - ); + expect(tracing.propagation.getActiveBaggage()).toBe(override); + const ctx = tracing.propagation.extract(tracing.context.active(), { + baggage: 'gen_ai.conversation.id=conv-from-header', + }); + await tracing.context.with(ctx, () => { + // mockReturnValue takes precedence over the default header parsing. + expect(tracing.propagation.getActiveBaggage()).toBe(override); + }); }); - it('returns undefined baggage when no baggage header is supplied', async () => { + it('returns undefined baggage when no baggage header is supplied to extract', async () => { const tracing = tracingServiceMock.mock(); - await tracing.withPropagatedContext({ traceparent: 'whatever' }, () => { - expect(tracing.getActiveBaggage()).toBeUndefined(); + const ctx = tracing.propagation.extract(tracing.context.active(), { + traceparent: 'whatever', + }); + await tracing.context.with(ctx, () => { + expect(tracing.propagation.getActiveBaggage()).toBeUndefined(); }); }); }); diff --git a/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts b/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts index 9f43c7149f..69125a7a2f 100644 --- a/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts +++ b/packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts @@ -22,22 +22,38 @@ import { TracingService, TracingServiceAttributeValue, TracingServiceBaggage, + TracingServiceContext, + TracingServiceContextAPI, + TracingServicePropagationAPI, TracingServiceSpan, TracingServiceSpanStatus, tracingServiceRef, } from '@backstage/backend-plugin-api/alpha'; import { tracingServiceFactory } from '@backstage/backend-defaults/alpha'; +// Internal context shape used by the mock. The opaque `TracingServiceContext` +// is just this object cast to the public type. +interface MockContext { + baggage?: TracingServiceBaggage; +} + +function toMockContext(ctx: TracingServiceContext): MockContext { + return ctx as unknown as MockContext; +} +function fromMockContext(ctx: MockContext): TracingServiceContext { + return ctx as unknown as TracingServiceContext; +} + // Parses the `baggage` header per the W3C Baggage member syntax, // dropping value properties (`;property=value`). This mirrors what // `propagation.extract` does in the real tracing service, just enough // for tests to assert end-to-end behaviour between propagated headers // and `getActiveBaggage()`. function parseBaggageHeader( - headers: Record, + carrier: Record, ): TracingServiceBaggage | undefined { let raw: string | undefined; - for (const [name, value] of Object.entries(headers)) { + for (const [name, value] of Object.entries(carrier)) { if (name.toLowerCase() !== 'baggage') continue; raw = Array.isArray(value) ? value[0] : value; break; @@ -72,27 +88,55 @@ export interface MockedTracingServiceSpan extends TracingServiceSpan { setStatus: jest.Mock; } +/** + * Jest-mocked counterpart of the `context` member on the + * `TracingService`. + * + * @alpha + */ +export interface MockedTracingServiceContextAPI + extends TracingServiceContextAPI { + active: jest.MockedFunction; + with: jest.MockedFunction; +} + +/** + * Jest-mocked counterpart of the `propagation` member on the + * `TracingService`. + * + * @alpha + */ +export interface MockedTracingServicePropagationAPI + extends TracingServicePropagationAPI { + extract: jest.MockedFunction; + getBaggage: jest.MockedFunction; + getActiveBaggage: jest.MockedFunction< + TracingServicePropagationAPI['getActiveBaggage'] + >; +} + /** * Mock for the `TracingService`. Captures every span created via * `startActiveSpan` so tests can assert on the options passed in and the * methods called on the span inside the callback. * - * By default, `withPropagatedContext` parses the `baggage` header (W3C - * Baggage syntax) out of the supplied headers and makes those entries - * available via `getActiveBaggage` for the duration of the wrapped - * callback. Other propagation headers (e.g. `traceparent`) are ignored. - * Tests that need fully custom baggage can still override - * `getActiveBaggage` via `mockReturnValue` / `mockImplementation`, which - * takes precedence over the default behaviour. + * By default, `propagation.extract` parses the `baggage` header (W3C + * Baggage syntax) out of the supplied carrier and stashes the entries + * on the returned context handle. `context.with` activates that handle + * for the duration of the wrapped callback so + * `propagation.getActiveBaggage` (and `propagation.getBaggage` on the + * supplied handle) returns those entries. Other propagation fields + * (e.g. `traceparent`) are ignored. Tests that need fully custom + * baggage can still override `propagation.getActiveBaggage` via + * `mockReturnValue` / `mockImplementation`, which takes precedence over + * the default behaviour. * * @alpha */ export interface TracingServiceMock extends TracingService { startActiveSpan: jest.MockedFunction; - withPropagatedContext: jest.MockedFunction< - TracingService['withPropagatedContext'] - >; - getActiveBaggage: jest.MockedFunction; + context: MockedTracingServiceContextAPI; + propagation: MockedTracingServicePropagationAPI; /** Spans created by `startActiveSpan` calls, in order. */ spans: MockedTracingServiceSpan[]; factory: ServiceFactory; @@ -113,35 +157,74 @@ export namespace tracingServiceMock { */ export const mock = (): TracingServiceMock => { const spans: MockedTracingServiceSpan[] = []; - const startActiveSpan = jest.fn(async (_name, fn, _options) => { - const span: MockedTracingServiceSpan = { - setAttribute: jest.fn(), - setStatus: jest.fn(), - }; - spans.push(span); - return await fn(span); - }) as TracingServiceMock['startActiveSpan']; + const startActiveSpan = jest.fn( + async ( + _name: string, + optionsOrFn: unknown, + maybeFn?: (span: MockedTracingServiceSpan) => unknown, + ) => { + const fn = ( + typeof optionsOrFn === 'function' ? optionsOrFn : maybeFn + ) as (span: MockedTracingServiceSpan) => unknown; + const span: MockedTracingServiceSpan = { + setAttribute: jest.fn(), + setStatus: jest.fn(), + }; + spans.push(span); + return await fn(span); + }, + ) as unknown as TracingServiceMock['startActiveSpan']; - const baggageStack: Array = []; - const withPropagatedContext = jest.fn(async (headers, fn) => { - baggageStack.push(parseBaggageHeader(headers)); + const contextStack: MockContext[] = [{}]; + + const active = jest.fn(() => + fromMockContext(contextStack[contextStack.length - 1]), + ) as MockedTracingServiceContextAPI['active']; + + const withFn = jest.fn(async (ctx, fn) => { + contextStack.push(toMockContext(ctx)); try { return await fn(); } finally { - baggageStack.pop(); + contextStack.pop(); } - }) as TracingServiceMock['withPropagatedContext']; - const getActiveBaggage = jest.fn( - () => baggageStack[baggageStack.length - 1], - ) as TracingServiceMock['getActiveBaggage']; + }) as MockedTracingServiceContextAPI['with']; - const service: TracingService = { - startActiveSpan, - withPropagatedContext, + const extract = jest.fn((ctx, carrier) => { + const baggage = parseBaggageHeader(carrier); + // Carry forward the parsed baggage; preserve any baggage already on the + // supplied handle if the carrier doesn't include one. + const base = toMockContext(ctx); + return fromMockContext({ baggage: baggage ?? base.baggage }); + }) as MockedTracingServicePropagationAPI['extract']; + + const getBaggage = jest.fn( + ctx => toMockContext(ctx).baggage, + ) as MockedTracingServicePropagationAPI['getBaggage']; + + const getActiveBaggage = jest.fn( + () => contextStack[contextStack.length - 1].baggage, + ) as MockedTracingServicePropagationAPI['getActiveBaggage']; + + const context: MockedTracingServiceContextAPI = { + active, + with: withFn, + }; + const propagation: MockedTracingServicePropagationAPI = { + extract, + getBaggage, getActiveBaggage, }; + const service: TracingService = { + startActiveSpan, + context, + propagation, + }; + return Object.assign(service as TracingServiceMock, { + context, + propagation, spans, factory: createServiceFactory({ service: tracingServiceRef, diff --git a/packages/backend-test-utils/src/alpha/services/index.ts b/packages/backend-test-utils/src/alpha/services/index.ts index c8861c86a4..e4c88f0b09 100644 --- a/packages/backend-test-utils/src/alpha/services/index.ts +++ b/packages/backend-test-utils/src/alpha/services/index.ts @@ -22,5 +22,7 @@ export { tracingServiceMock, type TracingServiceMock, type MockedTracingServiceSpan, + type MockedTracingServiceContextAPI, + type MockedTracingServicePropagationAPI, } from './TracingServiceMock'; export { type ServiceMock } from './alphaCreateServiceMock'; diff --git a/plugins/mcp-actions-backend/src/routers/createSseRouter.ts b/plugins/mcp-actions-backend/src/routers/createSseRouter.ts index d0f6411046..25418c69e9 100644 --- a/plugins/mcp-actions-backend/src/routers/createSseRouter.ts +++ b/plugins/mcp-actions-backend/src/routers/createSseRouter.ts @@ -68,7 +68,11 @@ export const createSseRouter = ({ const transport = transportsToSessionId.get(sessionId); if (transport) { - await tracing.withPropagatedContext(req.headers, () => + const ctx = tracing.propagation.extract( + tracing.context.active(), + req.headers, + ); + await tracing.context.with(ctx, () => transport.handlePostMessage(req, res, req.body), ); } else { diff --git a/plugins/mcp-actions-backend/src/routers/createStreamableRouter.ts b/plugins/mcp-actions-backend/src/routers/createStreamableRouter.ts index 1c71703ddf..61c9951b6d 100644 --- a/plugins/mcp-actions-backend/src/routers/createStreamableRouter.ts +++ b/plugins/mcp-actions-backend/src/routers/createStreamableRouter.ts @@ -77,7 +77,11 @@ export const createStreamableRouter = ({ }); await server.connect(transport); - await tracing.withPropagatedContext(req.headers, () => + const ctx = tracing.propagation.extract( + tracing.context.active(), + req.headers, + ); + await tracing.context.with(ctx, () => transport.handleRequest(req, res, req.body), ); diff --git a/plugins/mcp-actions-backend/src/services/McpService.test.ts b/plugins/mcp-actions-backend/src/services/McpService.test.ts index ae4a1770a4..7278bf276e 100644 --- a/plugins/mcp-actions-backend/src/services/McpService.test.ts +++ b/plugins/mcp-actions-backend/src/services/McpService.test.ts @@ -981,7 +981,7 @@ describe('McpService', () => { await invokeMockAction({ tracing, credentials }); expect(tracing.startActiveSpan).toHaveBeenCalledTimes(1); - const [name, , options] = tracing.startActiveSpan.mock.calls[0]; + const [name, options] = tracing.startActiveSpan.mock.calls[0]; expect(name).toBe('tools/call test.mock-action'); expect(options?.kind).toBe('server'); expect(options?.attributes).toEqual( @@ -1012,7 +1012,7 @@ describe('McpService', () => { it('includes gen_ai baggage entries as span attributes when present', async () => { const tracing = tracingServiceMock.mock(); - tracing.getActiveBaggage.mockReturnValue({ + tracing.propagation.getActiveBaggage.mockReturnValue({ getEntry: (key: string) => { const entries: Record = { 'gen_ai.conversation.id': { value: 'conv-123' }, @@ -1028,14 +1028,14 @@ describe('McpService', () => { await invokeMockAction({ tracing }); - const [, , options] = tracing.startActiveSpan.mock.calls[0]; + const [, options] = tracing.startActiveSpan.mock.calls[0]; expect(options?.attributes?.['gen_ai.conversation.id']).toBe('conv-123'); expect(options?.attributes?.['gen_ai.agent.id']).toBe('agent-456'); }); it('only forwards allowlisted baggage keys onto the span', async () => { const tracing = tracingServiceMock.mock(); - tracing.getActiveBaggage.mockReturnValue({ + tracing.propagation.getActiveBaggage.mockReturnValue({ getEntry: (key: string) => { const entries: Record = { 'gen_ai.conversation.id': { value: 'conv-123' }, @@ -1055,7 +1055,7 @@ describe('McpService', () => { await invokeMockAction({ tracing }); - const [, , options] = tracing.startActiveSpan.mock.calls[0]; + const [, options] = tracing.startActiveSpan.mock.calls[0]; expect(options?.attributes?.['gen_ai.conversation.id']).toBe('conv-123'); expect(options?.attributes).not.toHaveProperty('gen_ai.tool.call.result'); expect(options?.attributes).not.toHaveProperty('gen_ai.prompt'); @@ -1066,7 +1066,7 @@ describe('McpService', () => { const tracing = tracingServiceMock.mock(); await invokeMockAction({ tracing }); - const [, , options] = tracing.startActiveSpan.mock.calls[0]; + const [, options] = tracing.startActiveSpan.mock.calls[0]; expect(options?.attributes).not.toHaveProperty('gen_ai.conversation.id'); expect(options?.attributes).not.toHaveProperty('gen_ai.agent.id'); }); @@ -1075,7 +1075,7 @@ describe('McpService', () => { const tracing = tracingServiceMock.mock(); await invokeMockAction({ tracing, captureToolPayloads: true }); - const [, , options] = tracing.startActiveSpan.mock.calls[0]; + const [, options] = tracing.startActiveSpan.mock.calls[0]; expect(options?.attributes?.['gen_ai.tool.call.arguments']).toBe( JSON.stringify({ input: 'val' }), ); diff --git a/plugins/mcp-actions-backend/src/services/McpService.ts b/plugins/mcp-actions-backend/src/services/McpService.ts index 693f0b90a4..09685fed1d 100644 --- a/plugins/mcp-actions-backend/src/services/McpService.ts +++ b/plugins/mcp-actions-backend/src/services/McpService.ts @@ -57,7 +57,7 @@ const PROPAGATED_BAGGAGE_ATTRIBUTES: readonly string[] = [ function baggageAttributes( tracingService: TracingService, ): Record { - const baggage = tracingService.getActiveBaggage(); + const baggage = tracingService.propagation.getActiveBaggage(); if (!baggage) return {}; const attrs: Record = {}; for (const key of PROPAGATED_BAGGAGE_ATTRIBUTES) { @@ -189,6 +189,19 @@ export class McpService { try { return await this.tracingService.startActiveSpan( `tools/call ${params.name}`, + { + kind: 'server', + credentials, + attributes: { + ...baggageAttributes(this.tracingService), + 'mcp.method.name': 'tools/call', + 'gen_ai.tool.name': params.name, + 'gen_ai.operation.name': 'execute_tool', + ...(this.captureToolPayloads && { + 'gen_ai.tool.call.arguments': safeStringify(params.arguments), + }), + }, + }, async span => { const result = await handleErrors(async () => { const { actions: allActions } = await this.actions.list({ @@ -249,19 +262,6 @@ export class McpService { } return result; }, - { - kind: 'server', - credentials, - attributes: { - ...baggageAttributes(this.tracingService), - 'mcp.method.name': 'tools/call', - 'gen_ai.tool.name': params.name, - 'gen_ai.operation.name': 'execute_tool', - ...(this.captureToolPayloads && { - 'gen_ai.tool.call.arguments': safeStringify(params.arguments), - }), - }, - }, ); } catch (err) { errorType = err instanceof Error ? err.name : 'Error';