Propagate and forward gen_ai baggage from caller, if available
Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
@@ -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',
|
||||
|
||||
Reference in New Issue
Block a user