Propagate and forward gen_ai baggage from caller, if available

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
Eric Peterson
2026-05-12 13:36:17 +02:00
parent 6209065f00
commit d6c7805527
8 changed files with 273 additions and 18 deletions
@@ -89,6 +89,7 @@ export const mcpPlugin = createBackendPlugin({
httpAuth,
logger,
metrics,
tracing,
serverConfig,
});
@@ -105,6 +106,7 @@ export const mcpPlugin = createBackendPlugin({
const sseRouter = createSseRouter({
mcpService,
httpAuth,
tracing,
serverConfig,
});
@@ -113,6 +115,7 @@ export const mcpPlugin = createBackendPlugin({
httpAuth,
logger,
metrics,
tracing,
serverConfig,
});
@@ -18,6 +18,7 @@ import { Router } from 'express';
import { McpService } from '../services/McpService';
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
import { HttpAuthService } from '@backstage/backend-plugin-api';
import { TracingService } from '@backstage/backend-plugin-api/alpha';
import { McpServerConfig } from '../config';
/**
@@ -26,10 +27,12 @@ import { McpServerConfig } from '../config';
export const createSseRouter = ({
mcpService,
httpAuth,
tracing,
serverConfig,
}: {
mcpService: McpService;
httpAuth: HttpAuthService;
tracing: TracingService;
serverConfig?: McpServerConfig;
}): Router => {
const router = PromiseRouter();
@@ -65,7 +68,9 @@ export const createSseRouter = ({
const transport = transportsToSessionId.get(sessionId);
if (transport) {
await transport.handlePostMessage(req, res, req.body);
await tracing.withPropagatedContext(req.headers, () =>
transport.handlePostMessage(req, res, req.body),
);
} else {
res
.status(400)
@@ -21,7 +21,10 @@ import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/
import { LATEST_PROTOCOL_VERSION } from '@modelcontextprotocol/sdk/types.js';
import { HttpAuthService, LoggerService } from '@backstage/backend-plugin-api';
import { toError } from '@backstage/errors';
import { MetricsService } from '@backstage/backend-plugin-api/alpha';
import {
MetricsService,
TracingService,
} from '@backstage/backend-plugin-api/alpha';
import { bucketBoundaries, McpServerSessionAttributes } from '../metrics';
import { McpServerConfig } from '../config';
@@ -30,12 +33,14 @@ export const createStreamableRouter = ({
httpAuth,
logger,
metrics,
tracing,
serverConfig,
}: {
mcpService: McpService;
logger: LoggerService;
httpAuth: HttpAuthService;
metrics: MetricsService;
tracing: TracingService;
serverConfig?: McpServerConfig;
}): Router => {
const router = PromiseRouter();
@@ -72,7 +77,9 @@ export const createStreamableRouter = ({
});
await server.connect(transport);
await transport.handleRequest(req, res, req.body);
await tracing.withPropagatedContext(req.headers, () =>
transport.handleRequest(req, res, req.body),
);
res.on('close', () => {
transport.close();
@@ -1010,7 +1010,68 @@ describe('McpService', () => {
);
});
it('includes tool arguments in the span options and sets the result via setAttribute when captureToolPayloads is true', async () => {
it('includes gen_ai baggage entries as span attributes when present', async () => {
const tracing = tracingServiceMock.mock();
tracing.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' }],
],
});
await invokeMockAction({ tracing });
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({
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' }],
['gen_ai.prompt', { value: 'injected-prompt' }],
['gen_ai.user.message', { value: 'injected-user-message' }],
],
});
await invokeMockAction({ tracing });
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');
expect(options?.attributes).not.toHaveProperty('gen_ai.user.message');
});
it('omits gen_ai baggage attributes when no baggage is present', async () => {
const tracing = tracingServiceMock.mock();
await invokeMockAction({ tracing });
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');
});
it('includes tool arguments in the span options and sets the structured action output as the result attribute when captureToolPayloads is true', async () => {
const tracing = tracingServiceMock.mock();
await invokeMockAction({ tracing, captureToolPayloads: true });
@@ -1024,9 +1085,9 @@ describe('McpService', () => {
([key]) => key === 'gen_ai.tool.call.result',
);
expect(resultCall).toBeDefined();
const parsed = JSON.parse(resultCall![1] as string);
expect(parsed.content[0].type).toBe('text');
expect(parsed.content[0].text).toContain('"output": "test"');
// The recorded result should be the structured action output, not the
// CallToolResult envelope wrapping a fenced JSON block.
expect(JSON.parse(resultCall![1] as string)).toEqual({ output: 'test' });
});
it('sets error.type=tool_error and ERROR status on the span when the tool returns isError', async () => {
@@ -1049,6 +1110,7 @@ describe('McpService', () => {
actions: mockActionsRegistry,
metrics: metricsServiceMock.mock(),
tracingService: tracing,
captureToolPayloads: true,
});
const server = mcpService.getServer({
@@ -1083,6 +1145,12 @@ describe('McpService', () => {
code: 'error',
message: 'tool_error',
});
// The error is signalled via error.type + status; the result attribute
// should be omitted even when captureToolPayloads is enabled.
const resultCall = span.setAttribute.mock.calls.find(
([key]) => key === 'gen_ai.tool.call.result',
);
expect(resultCall).toBeUndefined();
});
});
});
@@ -43,6 +43,32 @@ 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[] = [
'gen_ai.agent.id',
'gen_ai.agent.name',
'gen_ai.conversation.id',
'gen_ai.provider.name',
'gen_ai.request.model',
];
function baggageAttributes(
tracingService: TracingService,
): Record<string, string> {
const baggage = tracingService.getActiveBaggage();
if (!baggage) return {};
const attrs: Record<string, string> = {};
for (const key of PROPAGATED_BAGGAGE_ATTRIBUTES) {
const entry = baggage.getEntry(key);
if (entry) {
attrs[key] = entry.value;
}
}
return attrs;
}
export class McpService {
private readonly actions: ActionsService;
private readonly namespacedToolNames: boolean;
@@ -189,6 +215,16 @@ export class McpService {
credentials,
});
// Record the structured action output directly rather than the
// CallToolResult envelope below, which wraps an already-
// stringified markdown-fenced JSON block.
if (this.captureToolPayloads) {
span.setAttribute(
'gen_ai.tool.call.result',
safeStringify(output),
);
}
return {
// todo(blam): unfortunately structuredContent is not supported by most clients yet.
// so the validation for the output happens in the default actions registry
@@ -207,12 +243,6 @@ export class McpService {
});
isError = !!(result as { isError?: boolean })?.isError;
if (this.captureToolPayloads) {
span.setAttribute(
'gen_ai.tool.call.result',
safeStringify(result),
);
}
if (isError) {
span.setAttribute('error.type', 'tool_error');
span.setStatus({ code: 'error', message: 'tool_error' });
@@ -223,6 +253,7 @@ export class McpService {
kind: 'server',
credentials,
attributes: {
...baggageAttributes(this.tracingService),
'mcp.method.name': 'tools/call',
'gen_ai.tool.name': params.name,
'gen_ai.operation.name': 'execute_tool',