Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,21 @@ test('Should record transactions for mcp handlers', async ({ baseURL }) => {
version: '1.0.0',
});

const initializeTransactionPromise = waitForTransaction('node-express-v5', transactionEvent => {
return transactionEvent.transaction === 'initialize';
});

await client.connect(transport);

await test.step('initialize handshake', async () => {
const initializeTransaction = await initializeTransactionPromise;
expect(initializeTransaction).toBeDefined();
expect(initializeTransaction.contexts?.trace?.op).toEqual('mcp.server');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.method.name']).toEqual('initialize');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.client.name']).toEqual('test-client');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.server.name']).toEqual('Echo');
});

await test.step('tool handler', async () => {
const postTransactionPromise = waitForTransaction('node-express-v5', transactionEvent => {
return transactionEvent.transaction === 'POST /messages';
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,21 @@ test('Should record transactions for mcp handlers', async ({ baseURL }) => {
version: '1.0.0',
});

const initializeTransactionPromise = waitForTransaction('node-express', transactionEvent => {
return transactionEvent.transaction === 'initialize';
});

await client.connect(transport);

await test.step('initialize handshake', async () => {
const initializeTransaction = await initializeTransactionPromise;
expect(initializeTransaction).toBeDefined();
expect(initializeTransaction.contexts?.trace?.op).toEqual('mcp.server');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.method.name']).toEqual('initialize');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.client.name']).toEqual('test-client');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.server.name']).toEqual('Echo');
});

await test.step('tool handler', async () => {
const postTransactionPromise = waitForTransaction('node-express', transactionEvent => {
return transactionEvent.transaction === 'POST /messages';
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,21 @@ test('Records transactions for mcp handlers', async ({ baseURL }) => {
version: '1.0.0',
});

const initializeTransactionPromise = waitForTransaction('tsx-express', transactionEvent => {
return transactionEvent.transaction === 'initialize';
});

await client.connect(transport);

await test.step('initialize handshake', async () => {
const initializeTransaction = await initializeTransactionPromise;
expect(initializeTransaction).toBeDefined();
expect(initializeTransaction.contexts?.trace?.op).toEqual('mcp.server');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.method.name']).toEqual('initialize');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.client.name']).toEqual('test-client');
expect(initializeTransaction.contexts?.trace?.data?.['mcp.server.name']).toEqual('Echo');
});

await test.step('tool handler', async () => {
const postTransactionPromise = waitForTransaction('tsx-express', transactionEvent => {
return transactionEvent.transaction === 'POST /messages';
Expand Down
20 changes: 17 additions & 3 deletions packages/core/src/integrations/mcp-server/correlation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@
import { getClient } from '../../currentScopes';
import { SPAN_STATUS_ERROR } from '../../tracing';
import type { Span } from '../../types-hoist/span';
import { MCP_PROTOCOL_VERSION_ATTRIBUTE } from './attributes';
import { filterMcpPiiFromSpanData } from './piiFiltering';
import { extractPromptResultAttributes, extractToolResultAttributes } from './resultExtraction';
import { buildServerAttributesFromInfo, extractSessionDataFromInitializeResponse } from './sessionExtraction';
import type { MCPTransport, RequestId, RequestSpanMapValue } from './types';

/**
Expand Down Expand Up @@ -51,18 +53,30 @@ export function storeSpanForRequest(transport: MCPTransport, requestId: RequestI
}

/**
* Completes span with tool results and cleans up correlation
* Completes span with results and cleans up correlation
* @param transport - MCP transport instance
* @param requestId - Request identifier
* @param result - Tool execution result for attribute extraction
* @param result - Execution result for attribute extraction
*/
export function completeSpanWithResults(transport: MCPTransport, requestId: RequestId, result: unknown): void {
const spanMap = getOrCreateSpanMap(transport);
const spanData = spanMap.get(requestId);
if (spanData) {
const { span, method } = spanData;

if (method === 'tools/call') {
if (method === 'initialize') {
const sessionData = extractSessionDataFromInitializeResponse(result);
const serverAttributes = buildServerAttributesFromInfo(sessionData.serverInfo);

const initAttributes: Record<string, string | number> = {
...serverAttributes,
};
if (sessionData.protocolVersion) {
initAttributes[MCP_PROTOCOL_VERSION_ATTRIBUTE] = sessionData.protocolVersion;
}

span.setAttributes(initAttributes);
} else if (method === 'tools/call') {
const rawToolAttributes = extractToolResultAttributes(result);
const client = getClient();
const sendDefaultPii = Boolean(client?.getOptions().sendDefaultPii);
Expand Down
42 changes: 42 additions & 0 deletions packages/core/src/integrations/mcp-server/sessionExtraction.ts
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,27 @@ export function getClientAttributes(transport: MCPTransport): Record<string, str
return attributes;
}

/**
* Build client attributes from PartyInfo directly
* @param clientInfo - Client party info
* @returns Client attributes for span instrumentation
*/
export function buildClientAttributesFromInfo(clientInfo?: PartyInfo): Record<string, string> {
const attributes: Record<string, string> = {};

if (clientInfo?.name) {
attributes['mcp.client.name'] = clientInfo.name;
}
if (clientInfo?.title) {
attributes['mcp.client.title'] = clientInfo.title;
}
if (clientInfo?.version) {
attributes['mcp.client.version'] = clientInfo.version;
}

return attributes;
}

/**
* Build server attributes from stored server info
* @param transport - MCP transport instance
Expand All @@ -127,6 +148,27 @@ export function getServerAttributes(transport: MCPTransport): Record<string, str
return attributes;
}

/**
* Build server attributes from PartyInfo directly
* @param serverInfo - Server party info
* @returns Server attributes for span instrumentation
*/
export function buildServerAttributesFromInfo(serverInfo?: PartyInfo): Record<string, string> {
const attributes: Record<string, string> = {};

if (serverInfo?.name) {
attributes[MCP_SERVER_NAME_ATTRIBUTE] = serverInfo.name;
}
if (serverInfo?.title) {
attributes[MCP_SERVER_TITLE_ATTRIBUTE] = serverInfo.title;
}
if (serverInfo?.version) {
attributes[MCP_SERVER_VERSION_ATTRIBUTE] = serverInfo.version;
}

return attributes;
}

/**
* Extracts client connection info from extra handler data
* @param extra - Extra handler data containing connection info
Expand Down
28 changes: 23 additions & 5 deletions packages/core/src/integrations/mcp-server/transport.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,16 +8,21 @@
import { getIsolationScope, withIsolationScope } from '../../currentScopes';
import { startInactiveSpan, withActiveSpan } from '../../tracing';
import { fill } from '../../utils/object';
import { MCP_PROTOCOL_VERSION_ATTRIBUTE } from './attributes';
import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation';
import { captureError } from './errorCapture';
import { extractSessionDataFromInitializeRequest, extractSessionDataFromInitializeResponse } from './sessionExtraction';
import {
buildClientAttributesFromInfo,
extractSessionDataFromInitializeRequest,
extractSessionDataFromInitializeResponse,
} from './sessionExtraction';
import {
cleanupSessionDataForTransport,
storeSessionDataForTransport,
updateSessionDataForTransport,
} from './sessionManagement';
import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans';
import type { ExtraHandlerData, MCPTransport } from './types';
import type { ExtraHandlerData, MCPTransport, SessionData } from './types';
import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse, isValidContentItem } from './validation';

/**
Expand All @@ -31,10 +36,13 @@ export function wrapTransportOnMessage(transport: MCPTransport): void {
fill(transport, 'onmessage', originalOnMessage => {
return function (this: MCPTransport, message: unknown, extra?: unknown) {
if (isJsonRpcRequest(message)) {
if (message.method === 'initialize') {
const isInitialize = message.method === 'initialize';
let initSessionData: SessionData | undefined;

if (isInitialize) {
try {
const sessionData = extractSessionDataFromInitializeRequest(message);
storeSessionDataForTransport(this, sessionData);
initSessionData = extractSessionDataFromInitializeRequest(message);
storeSessionDataForTransport(this, initSessionData);
} catch {
// noop
}
Expand All @@ -46,6 +54,16 @@ export function wrapTransportOnMessage(transport: MCPTransport): void {
const spanConfig = buildMcpServerSpanConfig(message, this, extra as ExtraHandlerData);
const span = startInactiveSpan(spanConfig);

// For initialize requests, add client info directly to span (works even for stateless transports)
if (isInitialize && initSessionData) {
span.setAttributes({
...buildClientAttributesFromInfo(initSessionData.clientInfo),
...(initSessionData.protocolVersion && {
[MCP_PROTOCOL_VERSION_ATTRIBUTE]: initSessionData.protocolVersion,
}),
});
}

storeSpanForRequest(this, message.id, span, message.method);

return withActiveSpan(span, () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -571,7 +571,7 @@ describe('MCP Server Transport Instrumentation', () => {

it('excludes sessionId when undefined', () => {
const transport = createMockTransport();
transport.sessionId = undefined;
transport.sessionId = '';
const attributes = buildTransportAttributes(transport);

expect(attributes['mcp.session.id']).toBeUndefined();
Expand All @@ -584,4 +584,75 @@ describe('MCP Server Transport Instrumentation', () => {
expect(attributes['mcp.session.id']).toBeUndefined();
});
});

describe('Initialize Span Attributes', () => {
it('should add client info to initialize span on request', async () => {
const mockMcpServer = createMockMcpServer();
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
const transport = createMockTransport();
transport.sessionId = '';

await wrappedMcpServer.connect(transport);

const mockSpan = { setAttributes: vi.fn(), end: vi.fn() };
startInactiveSpanSpy.mockReturnValue(mockSpan);

transport.onmessage?.(
{
jsonrpc: '2.0',
method: 'initialize',
id: 'init-1',
params: { protocolVersion: '2025-06-18', clientInfo: { name: 'test-client', version: '1.0.0' } },
},
{},
);

expect(mockSpan.setAttributes).toHaveBeenCalledWith(
expect.objectContaining({
'mcp.client.name': 'test-client',
'mcp.client.version': '1.0.0',
'mcp.protocol.version': '2025-06-18',
}),
);
});

it('should add server info to initialize span on response', async () => {
const mockMcpServer = createMockMcpServer();
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
const transport = createMockTransport();

await wrappedMcpServer.connect(transport);

const mockSpan = { setAttributes: vi.fn(), end: vi.fn() };
startInactiveSpanSpy.mockReturnValue(mockSpan as any);

transport.onmessage?.(
{
jsonrpc: '2.0',
method: 'initialize',
id: 'init-1',
params: { protocolVersion: '2025-06-18', clientInfo: { name: 'test-client', version: '1.0.0' } },
},
{},
);

await transport.send?.({
jsonrpc: '2.0',
id: 'init-1',
result: {
protocolVersion: '2025-06-18',
serverInfo: { name: 'test-server', version: '2.0.0' },
capabilities: {},
},
});

expect(mockSpan.setAttributes).toHaveBeenCalledWith(
expect.objectContaining({
'mcp.server.name': 'test-server',
'mcp.server.version': '2.0.0',
}),
);
expect(mockSpan.end).toHaveBeenCalled();
});
});
});
Loading