Skip to content

Commit ec3cb6f

Browse files
betegonbetegon
committed
refactor(mcp-server): span and attribute creation
1 parent fe2c865 commit ec3cb6f

File tree

3 files changed

+72
-127
lines changed

3 files changed

+72
-127
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,11 +71,11 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
7171

7272
transport.send = new Proxy(originalSend, {
7373
async apply(sendTarget, sendThisArg, sendArgs) {
74-
const [message, options] = sendArgs;
74+
const [message] = sendArgs;
7575

7676
// Instrument outgoing notifications (but not requests/responses)
7777
if (isJsonRpcNotification(message)) {
78-
return createMcpOutgoingNotificationSpan(message, transport, options as Record<string, unknown>, () => {
78+
return createMcpOutgoingNotificationSpan(message, transport, () => {
7979
return Reflect.apply(sendTarget, sendThisArg, sendArgs);
8080
});
8181
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,5 +110,16 @@ export interface ExtraHandlerData {
110110
};
111111
}
112112

113+
/**
114+
* Configuration for creating MCP spans
115+
*/
116+
export interface McpSpanConfig {
117+
type: 'request' | 'notification-incoming' | 'notification-outgoing';
118+
message: JsonRpcRequest | JsonRpcNotification;
119+
transport: MCPTransport;
120+
extra?: ExtraHandlerData;
121+
callback: () => unknown;
122+
}
123+
113124
export type SessionId = string;
114125
export type RequestId = string | number;

‎packages/core/src/utils/mcp-server/utils.ts

Lines changed: 59 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,9 @@ import {
2626
NETWORK_PROTOCOL_VERSION_ATTRIBUTE,
2727
NETWORK_TRANSPORT_ATTRIBUTE,
2828
} from './attributes';
29-
import type { ExtraHandlerData, JsonRpcNotification, JsonRpcRequest, MCPTransport } from './types';
29+
import type { ExtraHandlerData, JsonRpcNotification, JsonRpcRequest, McpSpanConfig, MCPTransport } from './types';
3030

31-
// =============================================================================
32-
// TYPE GUARDS
33-
// =============================================================================
34-
35-
/**
36-
*
37-
*/
31+
/** Validates if a message is a JSON-RPC request */
3832
export function isJsonRpcRequest(message: unknown): message is JsonRpcRequest {
3933
return (
4034
typeof message === 'object' &&
@@ -46,9 +40,7 @@ export function isJsonRpcRequest(message: unknown): message is JsonRpcRequest {
4640
);
4741
}
4842

49-
/**
50-
*
51-
*/
43+
/** Validates if a message is a JSON-RPC notification */
5244
export function isJsonRpcNotification(message: unknown): message is JsonRpcNotification {
5345
return (
5446
typeof message === 'object' &&
@@ -60,9 +52,7 @@ export function isJsonRpcNotification(message: unknown): message is JsonRpcNotif
6052
);
6153
}
6254

63-
/**
64-
*
65-
*/
55+
/** Extracts target info from method and params based on method type */
6656
export function validateMcpServerInstance(instance: unknown): boolean {
6757
if (
6858
typeof instance === 'object' &&
@@ -78,50 +68,39 @@ export function validateMcpServerInstance(instance: unknown): boolean {
7868
return false;
7969
}
8070

81-
// =============================================================================
82-
// ATTRIBUTE EXTRACTION
83-
// =============================================================================
71+
/** Extracts target info from method and params based on method type */
72+
function extractTargetInfo(method: string, params: Record<string, unknown>): {
73+
target?: string;
74+
attributes: Record<string, string>
75+
} {
76+
let target: string | undefined;
77+
let attributeKey: string | undefined;
8478

85-
/**
86-
*
87-
*/
88-
export function extractTarget(method: string, params: Record<string, unknown>): string | undefined {
8979
switch (method) {
9080
case 'tools/call':
91-
return typeof params?.name === 'string' ? params.name : undefined;
81+
target = typeof params?.name === 'string' ? params.name : undefined;
82+
attributeKey = 'mcp.tool.name';
83+
break;
9284
case 'resources/read':
9385
case 'resources/subscribe':
9486
case 'resources/unsubscribe':
95-
return typeof params?.uri === 'string' ? params.uri : undefined;
87+
target = typeof params?.uri === 'string' ? params.uri : undefined;
88+
attributeKey = 'mcp.resource.uri';
89+
break;
9690
case 'prompts/get':
97-
return typeof params?.name === 'string' ? params.name : undefined;
98-
default:
99-
return undefined;
91+
target = typeof params?.name === 'string' ? params.name : undefined;
92+
attributeKey = 'mcp.prompt.name';
93+
break;
10094
}
101-
}
10295

103-
/**
104-
*
105-
*/
106-
export function getTargetAttributes(method: string, target: string): Record<string, string> {
107-
switch (method) {
108-
case 'tools/call':
109-
return { 'mcp.tool.name': target };
110-
case 'resources/read':
111-
case 'resources/subscribe':
112-
case 'resources/unsubscribe':
113-
return { 'mcp.resource.uri': target };
114-
case 'prompts/get':
115-
return { 'mcp.prompt.name': target };
116-
default:
117-
return {};
118-
}
96+
return {
97+
target,
98+
attributes: target && attributeKey ? { [attributeKey]: target } : {}
99+
};
119100
}
120101

121-
/**
122-
*
123-
*/
124-
export function getRequestArguments(method: string, params: Record<string, unknown>): Record<string, string> {
102+
/** Extracts request arguments based on method type */
103+
function getRequestArguments(method: string, params: Record<string, unknown>): Record<string, string> {
125104
const args: Record<string, string> = {};
126105

127106
// Argument capture for different methods
@@ -153,14 +132,8 @@ export function getRequestArguments(method: string, params: Record<string, unkno
153132
return args;
154133
}
155134

156-
// =============================================================================
157-
// TRANSPORT DETECTION
158-
// =============================================================================
159-
160-
/**
161-
*
162-
*/
163-
export function getTransportTypes(transport: MCPTransport): { mcpTransport: string; networkTransport: string } {
135+
/** Extracts transport types based on transport constructor name */
136+
function getTransportTypes(transport: MCPTransport): { mcpTransport: string; networkTransport: string } {
164137
const transportName = transport.constructor?.name?.toLowerCase() || '';
165138

166139
if (transportName.includes('sse')) return { mcpTransport: 'sse', networkTransport: 'tcp' };
@@ -170,22 +143,8 @@ export function getTransportTypes(transport: MCPTransport): { mcpTransport: stri
170143
return { mcpTransport: 'http', networkTransport: 'tcp' };
171144
}
172145

173-
// =============================================================================
174-
// NOTIFICATION HANDLING
175-
// =============================================================================
176-
177-
/**
178-
* Get notification span name following OpenTelemetry conventions
179-
* For notifications, we use the method name directly as per JSON-RPC conventions
180-
*/
181-
export function getNotificationSpanName(method: string): string {
182-
return method;
183-
}
184-
185-
/**
186-
* Extract additional attributes for specific notification types
187-
*/
188-
export function getNotificationAttributes(
146+
/** Extracts additional attributes for specific notification types */
147+
function getNotificationAttributes(
189148
method: string,
190149
params: Record<string, unknown>,
191150
): Record<string, string | number> {
@@ -259,51 +218,14 @@ export function getNotificationAttributes(
259218
return attributes;
260219
}
261220

262-
// =============================================================================
263-
// CLIENT INFO EXTRACTION
264-
// =============================================================================
265-
266-
/**
267-
*
268-
*/
269-
export function extractClientAddress(extra: ExtraHandlerData): string | undefined {
270-
return (
271-
extra?.requestInfo?.remoteAddress ||
272-
extra?.clientAddress ||
273-
extra?.request?.ip ||
274-
extra?.request?.connection?.remoteAddress
275-
);
276-
}
277-
278-
/**
279-
*
280-
*/
281-
export function extractClientPort(extra: ExtraHandlerData): number | undefined {
282-
return extra?.requestInfo?.remotePort || extra?.clientPort || extra?.request?.connection?.remotePort;
283-
}
284-
285-
// =============================================================================
286-
// SPAN NAMING
287-
// =============================================================================
288221

289222
/**
290-
*
223+
* Creates a span name based on the method and target
291224
*/
292-
export function createSpanName(method: string, target?: string): string {
225+
function createSpanName(method: string, target?: string): string {
293226
return target ? `${method} ${target}` : method;
294227
}
295228

296-
// =============================================================================
297-
// UNIFIED SPAN BUILDER
298-
// =============================================================================
299-
300-
interface McpSpanConfig {
301-
type: 'request' | 'notification-incoming' | 'notification-outgoing';
302-
message: JsonRpcRequest | JsonRpcNotification;
303-
transport: MCPTransport;
304-
extra?: ExtraHandlerData;
305-
callback: () => unknown;
306-
}
307229

308230
/**
309231
* Unified builder for creating MCP spans
@@ -317,11 +239,11 @@ function createMcpSpan(config: McpSpanConfig): unknown {
317239
// Determine span name based on type and OTEL conventions
318240
let spanName: string;
319241
if (type === 'request') {
320-
const target = extractTarget(method, params || {});
321-
spanName = createSpanName(method, target);
242+
const targetInfo = extractTargetInfo(method, params || {});
243+
spanName = createSpanName(method, targetInfo.target);
322244
} else {
323-
// For notifications, use method name directly (OTEL convention)
324-
spanName = getNotificationSpanName(method);
245+
// For notifications, use method name directly per OpenTelemetry conventions
246+
spanName = method;
325247
}
326248

327249
// Build attributes
@@ -354,14 +276,13 @@ function buildTransportAttributes(
354276
extra?: ExtraHandlerData,
355277
): Record<string, string | number> {
356278
const sessionId = transport.sessionId;
357-
const clientAddress = extra ? extractClientAddress(extra) : undefined;
358-
const clientPort = extra ? extractClientPort(extra) : undefined;
279+
const clientInfo = extra ? extractClientInfo(extra) : {};
359280
const { mcpTransport, networkTransport } = getTransportTypes(transport);
360281

361282
return {
362283
...(sessionId && { [MCP_SESSION_ID_ATTRIBUTE]: sessionId }),
363-
...(clientAddress && { [CLIENT_ADDRESS_ATTRIBUTE]: clientAddress }),
364-
...(clientPort && { [CLIENT_PORT_ATTRIBUTE]: clientPort }),
284+
...(clientInfo.address && { [CLIENT_ADDRESS_ATTRIBUTE]: clientInfo.address }),
285+
...(clientInfo.port && { [CLIENT_PORT_ATTRIBUTE]: clientInfo.port }),
365286
[MCP_TRANSPORT_ATTRIBUTE]: mcpTransport,
366287
[NETWORK_TRANSPORT_ATTRIBUTE]: networkTransport,
367288
[NETWORK_PROTOCOL_VERSION_ATTRIBUTE]: '2.0',
@@ -378,11 +299,11 @@ function buildTypeSpecificAttributes(
378299
): Record<string, string | number> {
379300
if (type === 'request') {
380301
const request = message as JsonRpcRequest;
381-
const target = extractTarget(request.method, params || {});
302+
const targetInfo = extractTargetInfo(request.method, params || {});
382303

383304
return {
384305
...(request.id !== undefined && { [MCP_REQUEST_ID_ATTRIBUTE]: String(request.id) }),
385-
...(target && getTargetAttributes(request.method, target)),
306+
...targetInfo.attributes,
386307
...getRequestArguments(request.method, params || {}),
387308
};
388309
}
@@ -421,10 +342,6 @@ function buildSentryAttributes(type: McpSpanConfig['type']): Record<string, stri
421342
};
422343
}
423344

424-
// =============================================================================
425-
// PUBLIC API - SIMPLIFIED SPAN CREATION FUNCTIONS
426-
// =============================================================================
427-
428345
/**
429346
* Creates a span for MCP server request handling
430347
*/
@@ -467,7 +384,6 @@ export function createMcpNotificationSpan(
467384
export function createMcpOutgoingNotificationSpan(
468385
jsonRpcMessage: JsonRpcNotification,
469386
transport: MCPTransport,
470-
options: Record<string, unknown>,
471387
callback: () => unknown,
472388
): unknown {
473389
return createMcpSpan({
@@ -477,3 +393,21 @@ export function createMcpOutgoingNotificationSpan(
477393
callback,
478394
});
479395
}
396+
397+
/**
398+
* Combine the two extraction functions into one
399+
*/
400+
function extractClientInfo(extra: ExtraHandlerData): {
401+
address?: string;
402+
port?: number
403+
} {
404+
return {
405+
address: extra?.requestInfo?.remoteAddress ||
406+
extra?.clientAddress ||
407+
extra?.request?.ip ||
408+
extra?.request?.connection?.remoteAddress,
409+
port: extra?.requestInfo?.remotePort ||
410+
extra?.clientPort ||
411+
extra?.request?.connection?.remotePort
412+
};
413+
}

0 commit comments

Comments
 (0)