Skip to content

Commit 5a6f4ab

Browse files
betegonbetegon
committed
Get and persist session data + refactor to fit the max 300 lines file requirement
1 parent 64bfeda commit 5a6f4ab

File tree

8 files changed

+425
-210
lines changed

8 files changed

+425
-210
lines changed

‎packages/core/src/integrations/mcp-server/attributeExtraction.ts

Lines changed: 150 additions & 187 deletions
Large diffs are not rendered by default.

‎packages/core/src/integrations/mcp-server/attributes.ts

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,35 @@ export const MCP_SESSION_ID_ATTRIBUTE = 'mcp.session.id';
2121
/** Transport method used for MCP communication */
2222
export const MCP_TRANSPORT_ATTRIBUTE = 'mcp.transport';
2323

24+
// =============================================================================
25+
// CLIENT ATTRIBUTES
26+
// =============================================================================
27+
28+
/** Name of the MCP client application */
29+
export const MCP_CLIENT_NAME_ATTRIBUTE = 'mcp.client.name';
30+
31+
/** Display title of the MCP client application */
32+
export const MCP_CLIENT_TITLE_ATTRIBUTE = 'mcp.client.title';
33+
34+
/** Version of the MCP client application */
35+
export const MCP_CLIENT_VERSION_ATTRIBUTE = 'mcp.client.version';
36+
37+
// =============================================================================
38+
// SERVER ATTRIBUTES
39+
// =============================================================================
40+
41+
/** Name of the MCP server application */
42+
export const MCP_SERVER_NAME_ATTRIBUTE = 'mcp.server.name';
43+
44+
/** Display title of the MCP server application */
45+
export const MCP_SERVER_TITLE_ATTRIBUTE = 'mcp.server.title';
46+
47+
/** Version of the MCP server application */
48+
export const MCP_SERVER_VERSION_ATTRIBUTE = 'mcp.server.version';
49+
50+
/** MCP protocol version used in the session */
51+
export const MCP_PROTOCOL_VERSION_ATTRIBUTE = 'mcp.protocol.version';
52+
2453
// =============================================================================
2554
// METHOD-SPECIFIC ATTRIBUTES
2655
// =============================================================================

‎packages/core/src/integrations/mcp-server/methodConfig.ts

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
/**
2+
* Method configuration and request processing for MCP server instrumentation
3+
*/
4+
5+
import {
6+
MCP_PROMPT_NAME_ATTRIBUTE,
7+
MCP_REQUEST_ARGUMENT,
8+
MCP_RESOURCE_URI_ATTRIBUTE,
9+
MCP_TOOL_NAME_ATTRIBUTE,
10+
} from './attributes';
11+
import type { MethodConfig } from './types';
12+
13+
/**
14+
* Configuration for MCP methods to extract targets and arguments
15+
* @internal Maps method names to their extraction configuration
16+
*/
17+
const METHOD_CONFIGS: Record<string, MethodConfig> = {
18+
'tools/call': {
19+
targetField: 'name',
20+
targetAttribute: MCP_TOOL_NAME_ATTRIBUTE,
21+
captureArguments: true,
22+
argumentsField: 'arguments',
23+
},
24+
'resources/read': {
25+
targetField: 'uri',
26+
targetAttribute: MCP_RESOURCE_URI_ATTRIBUTE,
27+
captureUri: true,
28+
},
29+
'resources/subscribe': {
30+
targetField: 'uri',
31+
targetAttribute: MCP_RESOURCE_URI_ATTRIBUTE,
32+
},
33+
'resources/unsubscribe': {
34+
targetField: 'uri',
35+
targetAttribute: MCP_RESOURCE_URI_ATTRIBUTE,
36+
},
37+
'prompts/get': {
38+
targetField: 'name',
39+
targetAttribute: MCP_PROMPT_NAME_ATTRIBUTE,
40+
captureName: true,
41+
captureArguments: true,
42+
argumentsField: 'arguments',
43+
},
44+
};
45+
46+
/**
47+
* Extracts target info from method and params based on method type
48+
* @param method - MCP method name
49+
* @param params - Method parameters
50+
* @returns Target name and attributes for span instrumentation
51+
*/
52+
export function extractTargetInfo(
53+
method: string,
54+
params: Record<string, unknown>,
55+
): {
56+
target?: string;
57+
attributes: Record<string, string>;
58+
} {
59+
const config = METHOD_CONFIGS[method as keyof typeof METHOD_CONFIGS];
60+
if (!config) {
61+
return { attributes: {} };
62+
}
63+
64+
const target =
65+
config.targetField && typeof params?.[config.targetField] === 'string'
66+
? (params[config.targetField] as string)
67+
: undefined;
68+
69+
return {
70+
target,
71+
attributes: target && config.targetAttribute ? { [config.targetAttribute]: target } : {},
72+
};
73+
}
74+
75+
/**
76+
* Extracts request arguments based on method type
77+
* @param method - MCP method name
78+
* @param params - Method parameters
79+
* @returns Arguments as span attributes with mcp.request.argument prefix
80+
*/
81+
export function getRequestArguments(method: string, params: Record<string, unknown>): Record<string, string> {
82+
const args: Record<string, string> = {};
83+
const config = METHOD_CONFIGS[method as keyof typeof METHOD_CONFIGS];
84+
85+
if (!config) {
86+
return args;
87+
}
88+
89+
if (config.captureArguments && config.argumentsField && params?.[config.argumentsField]) {
90+
const argumentsObj = params[config.argumentsField];
91+
if (typeof argumentsObj === 'object' && argumentsObj !== null) {
92+
for (const [key, value] of Object.entries(argumentsObj as Record<string, unknown>)) {
93+
args[`${MCP_REQUEST_ARGUMENT}.${key.toLowerCase()}`] = JSON.stringify(value);
94+
}
95+
}
96+
}
97+
98+
if (config.captureUri && params?.uri) {
99+
args[`${MCP_REQUEST_ARGUMENT}.uri`] = JSON.stringify(params.uri);
100+
}
101+
102+
if (config.captureName && params?.name) {
103+
args[`${MCP_REQUEST_ARGUMENT}.name`] = JSON.stringify(params.name);
104+
}
105+
106+
return args;
107+
}

‎packages/core/src/integrations/mcp-server/sessionManagement.ts

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
/**
2+
* Session data management for MCP server instrumentation
3+
*/
4+
5+
import type { MCPTransport, PartyInfo, SessionData } from './types';
6+
7+
/**
8+
* Transport-scoped session data storage (only for transports with sessionId)
9+
* @internal Maps transport instances to session-level data
10+
*/
11+
const transportToSessionData = new WeakMap<MCPTransport, SessionData>();
12+
13+
/**
14+
* Stores session data for a transport with sessionId
15+
* @param transport - MCP transport instance
16+
* @param sessionData - Session data to store
17+
*/
18+
export function storeSessionDataForTransport(transport: MCPTransport, sessionData: SessionData): void {
19+
if (transport.sessionId) transportToSessionData.set(transport, sessionData);
20+
}
21+
22+
/**
23+
* Updates session data for a transport with sessionId (merges with existing data)
24+
* @param transport - MCP transport instance
25+
* @param partialSessionData - Partial session data to merge with existing data
26+
*/
27+
export function updateSessionDataForTransport(transport: MCPTransport, partialSessionData: Partial<SessionData>): void {
28+
if (transport.sessionId) {
29+
const existingData = transportToSessionData.get(transport) || {};
30+
transportToSessionData.set(transport, { ...existingData, ...partialSessionData });
31+
}
32+
}
33+
34+
/**
35+
* Retrieves client information for a transport
36+
* @param transport - MCP transport instance
37+
* @returns Client information if available
38+
*/
39+
export function getClientInfoForTransport(transport: MCPTransport): PartyInfo | undefined {
40+
return transportToSessionData.get(transport)?.clientInfo;
41+
}
42+
43+
/**
44+
* Retrieves protocol version for a transport
45+
* @param transport - MCP transport instance
46+
* @returns Protocol version if available
47+
*/
48+
export function getProtocolVersionForTransport(transport: MCPTransport): string | undefined {
49+
return transportToSessionData.get(transport)?.protocolVersion;
50+
}
51+
52+
/**
53+
* Retrieves full session data for a transport
54+
* @param transport - MCP transport instance
55+
* @returns Complete session data if available
56+
*/
57+
export function getSessionDataForTransport(transport: MCPTransport): SessionData | undefined {
58+
return transportToSessionData.get(transport);
59+
}
60+
61+
/**
62+
* Cleans up session data for a specific transport (when that transport closes)
63+
* @param transport - MCP transport instance
64+
*/
65+
export function cleanupSessionDataForTransport(transport: MCPTransport): void {
66+
transportToSessionData.delete(transport);
67+
}

‎packages/core/src/integrations/mcp-server/spans.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import {
1212
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
1313
} from '../../semanticAttributes';
1414
import { startSpan } from '../../tracing';
15-
import { buildTransportAttributes, buildTypeSpecificAttributes, extractTargetInfo } from './attributeExtraction';
15+
import { buildTransportAttributes, buildTypeSpecificAttributes } from './attributeExtraction';
1616
import {
1717
MCP_FUNCTION_ORIGIN_VALUE,
1818
MCP_METHOD_NAME_ATTRIBUTE,
@@ -22,6 +22,7 @@ import {
2222
MCP_ROUTE_SOURCE_VALUE,
2323
MCP_SERVER_OP_VALUE,
2424
} from './attributes';
25+
import { extractTargetInfo } from './methodConfig';
2526
import { filterMcpPiiFromSpanData } from './piiFiltering';
2627
import type { ExtraHandlerData, JsonRpcNotification, JsonRpcRequest, McpSpanConfig, MCPTransport } from './types';
2728

‎packages/core/src/integrations/mcp-server/transport.ts

Lines changed: 48 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,51 +8,71 @@
88
import { getIsolationScope, withIsolationScope } from '../../currentScopes';
99
import { startInactiveSpan, withActiveSpan } from '../../tracing';
1010
import { fill } from '../../utils/object';
11+
import {
12+
extractSessionDataFromInitializeRequest,
13+
extractSessionDataFromInitializeResponse,
14+
} from './attributeExtraction';
1115
import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation';
1216
import { captureError } from './errorCapture';
17+
import {
18+
cleanupSessionDataForTransport,
19+
storeSessionDataForTransport,
20+
updateSessionDataForTransport,
21+
} from './sessionManagement';
1322
import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans';
1423
import type { ExtraHandlerData, MCPTransport } from './types';
1524
import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse } from './validation';
1625

1726
/**
18-
* Wraps transport.onmessage to create spans for incoming messages
27+
* Wraps transport.onmessage to create spans for incoming messages.
28+
* For "initialize" requests, extracts and stores client info and protocol version
29+
* in the session data for the transport.
1930
* @param transport - MCP transport instance to wrap
2031
*/
2132
export function wrapTransportOnMessage(transport: MCPTransport): void {
2233
if (transport.onmessage) {
2334
fill(transport, 'onmessage', originalOnMessage => {
24-
return function (this: MCPTransport, jsonRpcMessage: unknown, extra?: unknown) {
25-
if (isJsonRpcRequest(jsonRpcMessage)) {
26-
const messageTyped = jsonRpcMessage as { method: string; id: string | number };
35+
return function (this: MCPTransport, message: unknown, extra?: unknown) {
36+
if (isJsonRpcRequest(message)) {
37+
if (message.method === 'initialize') {
38+
try {
39+
const sessionData = extractSessionDataFromInitializeRequest(message);
40+
storeSessionDataForTransport(this, sessionData);
41+
} catch {
42+
// noop
43+
}
44+
}
2745

2846
const isolationScope = getIsolationScope().clone();
2947

3048
return withIsolationScope(isolationScope, () => {
31-
const spanConfig = buildMcpServerSpanConfig(jsonRpcMessage, this, extra as ExtraHandlerData);
49+
const spanConfig = buildMcpServerSpanConfig(message, this, extra as ExtraHandlerData);
3250
const span = startInactiveSpan(spanConfig);
3351

34-
storeSpanForRequest(this, messageTyped.id, span, messageTyped.method);
52+
storeSpanForRequest(this, message.id, span, message.method);
3553

3654
return withActiveSpan(span, () => {
37-
return (originalOnMessage as (...args: unknown[]) => unknown).call(this, jsonRpcMessage, extra);
55+
return (originalOnMessage as (...args: unknown[]) => unknown).call(this, message, extra);
3856
});
3957
});
4058
}
4159

42-
if (isJsonRpcNotification(jsonRpcMessage)) {
43-
return createMcpNotificationSpan(jsonRpcMessage, this, extra as ExtraHandlerData, () => {
44-
return (originalOnMessage as (...args: unknown[]) => unknown).call(this, jsonRpcMessage, extra);
60+
if (isJsonRpcNotification(message)) {
61+
return createMcpNotificationSpan(message, this, extra as ExtraHandlerData, () => {
62+
return (originalOnMessage as (...args: unknown[]) => unknown).call(this, message, extra);
4563
});
4664
}
4765

48-
return (originalOnMessage as (...args: unknown[]) => unknown).call(this, jsonRpcMessage, extra);
66+
return (originalOnMessage as (...args: unknown[]) => unknown).call(this, message, extra);
4967
};
5068
});
5169
}
5270
}
5371

5472
/**
55-
* Wraps transport.send to handle outgoing messages and response correlation
73+
* Wraps transport.send to handle outgoing messages and response correlation.
74+
* For "initialize" responses, extracts and stores protocol version and server info
75+
* in the session data for the transport.
5676
* @param transport - MCP transport instance to wrap
5777
*/
5878
export function wrapTransportSend(transport: MCPTransport): void {
@@ -66,14 +86,24 @@ export function wrapTransportSend(transport: MCPTransport): void {
6686
}
6787

6888
if (isJsonRpcResponse(message)) {
69-
const messageTyped = message as { id: string | number; result?: unknown; error?: unknown };
89+
if (message.id !== null && message.id !== undefined) {
90+
if (message.error) {
91+
captureJsonRpcErrorResponse(message.error);
92+
}
7093

71-
if (messageTyped.id !== null && messageTyped.id !== undefined) {
72-
if (messageTyped.error) {
73-
captureJsonRpcErrorResponse(messageTyped.error);
94+
if (message.result && typeof message.result === 'object') {
95+
const result = message.result as Record<string, unknown>;
96+
if (result.protocolVersion || result.serverInfo) {
97+
try {
98+
const serverData = extractSessionDataFromInitializeResponse(message.result);
99+
updateSessionDataForTransport(this, serverData);
100+
} catch {
101+
// noop
102+
}
103+
}
74104
}
75105

76-
completeSpanWithResults(this, messageTyped.id, messageTyped.result);
106+
completeSpanWithResults(this, message.id, message.result);
77107
}
78108
}
79109

@@ -92,7 +122,7 @@ export function wrapTransportOnClose(transport: MCPTransport): void {
92122
fill(transport, 'onclose', originalOnClose => {
93123
return function (this: MCPTransport, ...args: unknown[]) {
94124
cleanupPendingSpansForTransport(this);
95-
125+
cleanupSessionDataForTransport(this);
96126
return (originalOnClose as (...args: unknown[]) => unknown).call(this, ...args);
97127
};
98128
});

‎packages/core/src/integrations/mcp-server/types.ts

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,3 +163,23 @@ export type McpErrorType =
163163
| 'protocol'
164164
| 'validation'
165165
| 'timeout';
166+
167+
/**
168+
* Party (client/server) information extracted from MCP initialize requests
169+
* @internal
170+
*/
171+
export type PartyInfo = {
172+
name?: string;
173+
title?: string;
174+
version?: string;
175+
};
176+
177+
/**
178+
* Session-level data collected from various MCP messages
179+
* @internal
180+
*/
181+
export type SessionData = {
182+
clientInfo?: PartyInfo;
183+
protocolVersion?: string;
184+
serverInfo?: PartyInfo;
185+
};

‎packages/core/src/integrations/mcp-server/validation.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
import { DEBUG_BUILD } from '../../debug-build';
88
import { debug } from '../../utils/debug-logger';
9-
import type { JsonRpcNotification, JsonRpcRequest } from './types';
9+
import type { JsonRpcNotification, JsonRpcRequest, JsonRpcResponse } from './types';
1010

1111
/**
1212
* Validates if a message is a JSON-RPC request
@@ -45,9 +45,7 @@ export function isJsonRpcNotification(message: unknown): message is JsonRpcNotif
4545
* @param message - Message to validate
4646
* @returns True if message is a JSON-RPC response
4747
*/
48-
export function isJsonRpcResponse(
49-
message: unknown,
50-
): message is { jsonrpc: '2.0'; id: string | number | null; result?: unknown; error?: unknown } {
48+
export function isJsonRpcResponse(message: unknown): message is JsonRpcResponse {
5149
return (
5250
typeof message === 'object' &&
5351
message !== null &&

0 commit comments

Comments
 (0)