Merge pull request #34089 from backstage/otel/mcp-tools-call

feat: Instrument MCP tool calls with semantically appropriate span

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<TInputSchema, TOutputSchema>): void;
 }
 
+// @alpha
+export interface MockedTracingServiceContextAPI
+  extends TracingServiceContextAPI {
+  // (undocumented)
+  active: jest.MockedFunction<TracingServiceContextAPI['active']>;
+  // (undocumented)
+  with: jest.MockedFunction<TracingServiceContextAPI['with']>;
+}
+
+// @alpha
+export interface MockedTracingServicePropagationAPI
+  extends TracingServicePropagationAPI {
+  // (undocumented)
+  extract: jest.MockedFunction<TracingServicePropagationAPI['extract']>;
+  // (undocumented)
+  getActiveBaggage: jest.MockedFunction<
+    TracingServicePropagationAPI['getActiveBaggage']
+  >;
+  // (undocumented)
+  getBaggage: jest.MockedFunction<TracingServicePropagationAPI['getBaggage']>;
+}
+
 // @alpha
 export interface MockedTracingServiceSpan extends TracingServiceSpan {
   // (undocumented)
@@ -106,8 +130,12 @@ export type ServiceMock<TService> = {
 
 // @alpha
 export interface TracingServiceMock extends TracingService {
+  // (undocumented)
+  context: MockedTracingServiceContextAPI;
   // (undocumented)
   factory: ServiceFactory<TracingService>;
+  // (undocumented)
+  propagation: MockedTracingServicePropagationAPI;
   spans: MockedTracingServiceSpan[];
   // (undocumented)
   startActiveSpan: jest.MockedFunction<TracingService['startActiveSpan']>;

packages/backend-test-utils/src/alpha/services/TracingServiceMock.test.ts

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2026 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { tracingServiceMock } from './TracingServiceMock';
+
+describe('tracingServiceMock', () => {
+  it('parses the baggage header via propagation.extract and exposes it via getActiveBaggage inside context.with', async () => {
+    const tracing = tracingServiceMock.mock();
+
+    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)?.getAllEntries()).toEqual([
+      ['gen_ai.conversation.id', { value: 'conv-123' }],
+      ['gen_ai.agent.id', { value: 'agent-456' }],
+    ]);
+
+    const seen = await tracing.context.with(ctx, () =>
+      tracing.propagation
+        .getActiveBaggage()
+        ?.getAllEntries()
+        .map(([k, v]) => [k, v.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();
+  });
+
+  it('honours mockReturnValue overrides for getActiveBaggage', async () => {
+    const tracing = tracingServiceMock.mock();
+    const override = {
+      getAllEntries: () =>
+        [['gen_ai.conversation.id', { value: 'override' }]] as Array<
+          [string, { value: string }]
+        >,
+    };
+    tracing.propagation.getActiveBaggage.mockReturnValue(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 to extract', async () => {
+    const tracing = tracingServiceMock.mock();
+    const ctx = tracing.propagation.extract(tracing.context.active(), {
+      traceparent: 'whatever',
+    });
+    await tracing.context.with(ctx, () => {
+      expect(tracing.propagation.getActiveBaggage()).toBeUndefined();
+    });
+  });
+});

packages/backend-test-utils/src/alpha/services/TracingServiceMock.ts

@@ -21,12 +21,62 @@ import {
 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(
+  carrier: Record<string, string | string[] | undefined>,
+): TracingServiceBaggage | undefined {
+  let raw: string | undefined;
+  for (const [name, value] of Object.entries(carrier)) {
+    if (name.toLowerCase() !== 'baggage') continue;
+    raw = Array.isArray(value) ? value[0] : value;
+    break;
+  }
+  if (!raw) return undefined;
+
+  const entries = new Map<string, { value: string }>();
+  for (const segment of raw.split(',')) {
+    const [pair] = segment.split(';');
+    const eqIdx = pair.indexOf('=');
+    if (eqIdx === -1) continue;
+    const key = decodeURIComponent(pair.slice(0, eqIdx).trim());
+    const value = decodeURIComponent(pair.slice(eqIdx + 1).trim());
+    if (!key) continue;
+    entries.set(key, { value });
+  }
+  if (entries.size === 0) return undefined;
+
+  return {
+    getAllEntries: () => Array.from(entries.entries()),
+  };
+}
+
 /**
  * A jest-mocked span captured by {@link TracingServiceMock}.
  *
@@ -37,15 +87,61 @@ export interface MockedTracingServiceSpan extends TracingServiceSpan {
   setStatus: jest.Mock<void, [TracingServiceSpanStatus]>;
 }
 
+/**
+ * Jest-mocked counterpart of the `context` member on the
+ * `TracingService`.
+ *
+ * @alpha
+ */
+export interface MockedTracingServiceContextAPI
+  extends TracingServiceContextAPI {
+  active: jest.MockedFunction<TracingServiceContextAPI['active']>;
+  with: jest.MockedFunction<TracingServiceContextAPI['with']>;
+}
+
+/**
+ * Jest-mocked counterpart of the `propagation` member on the
+ * `TracingService`.
+ *
+ * @alpha
+ */
+export interface MockedTracingServicePropagationAPI
+  extends TracingServicePropagationAPI {
+  extract: jest.MockedFunction<TracingServicePropagationAPI['extract']>;
+  getBaggage: jest.MockedFunction<TracingServicePropagationAPI['getBaggage']>;
+  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, `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.
+ *
+ * Unlike the real `DefaultTracingService`, the mock's `startActiveSpan`
+ * does **not** resolve `options.credentials` from `options.request` via
+ * `httpAuth`. Tests that need principal-derived span attributes should
+ * supply `options.credentials` directly on the span options, or assert
+ * on the raw `options` captured by `startActiveSpan.mock.calls`.
+ *
  * @alpha
  */
 export interface TracingServiceMock extends TracingService {
   startActiveSpan: jest.MockedFunction<TracingService['startActiveSpan']>;
+  context: MockedTracingServiceContextAPI;
+  propagation: MockedTracingServicePropagationAPI;
   /** Spans created by `startActiveSpan` calls, in order. */
   spans: MockedTracingServiceSpan[];
   factory: ServiceFactory<TracingService>;
@@ -66,18 +162,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 service: TracingService = { startActiveSpan };
+    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 {
+        contextStack.pop();
+      }
+    }) as MockedTracingServiceContextAPI['with'];
+
+    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,

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';