Skip to content

Commit de1b87f

Browse files
betegonbetegon
committed
Refactor MCP server integration. improve attribute extraction and span handling functions. Update type definitions and separate method wrapping for transport handlers.
1 parent 4973a60 commit de1b87f

File tree

9 files changed

+481
-437
lines changed

9 files changed

+481
-437
lines changed

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

Lines changed: 123 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,14 @@ import {
1515
NETWORK_PROTOCOL_VERSION_ATTRIBUTE,
1616
NETWORK_TRANSPORT_ATTRIBUTE,
1717
} from './attributes';
18-
import type { ExtraHandlerData, JsonRpcNotification, JsonRpcRequest, McpSpanType,MCPTransport, MethodConfig } from './types';
18+
import type {
19+
ExtraHandlerData,
20+
JsonRpcNotification,
21+
JsonRpcRequest,
22+
McpSpanType,
23+
MCPTransport,
24+
MethodConfig,
25+
} from './types';
1926

2027
/** Configuration for MCP methods to extract targets and arguments */
2128
const METHOD_CONFIGS: Record<string, MethodConfig> = {
@@ -48,22 +55,26 @@ const METHOD_CONFIGS: Record<string, MethodConfig> = {
4855
};
4956

5057
/** Extracts target info from method and params based on method type */
51-
export function extractTargetInfo(method: string, params: Record<string, unknown>): {
58+
export function extractTargetInfo(
59+
method: string,
60+
params: Record<string, unknown>,
61+
): {
5262
target?: string;
53-
attributes: Record<string, string>
63+
attributes: Record<string, string>;
5464
} {
5565
const config = METHOD_CONFIGS[method as keyof typeof METHOD_CONFIGS];
5666
if (!config) {
5767
return { attributes: {} };
5868
}
5969

60-
const target = config.targetField && typeof params?.[config.targetField] === 'string'
61-
? params[config.targetField] as string
62-
: undefined;
70+
const target =
71+
config.targetField && typeof params?.[config.targetField] === 'string'
72+
? (params[config.targetField] as string)
73+
: undefined;
6374

6475
return {
6576
target,
66-
attributes: target && config.targetAttribute ? { [config.targetAttribute]: target } : {}
77+
attributes: target && config.targetAttribute ? { [config.targetAttribute]: target } : {},
6778
};
6879
}
6980

@@ -197,16 +208,15 @@ export function getNotificationAttributes(
197208
/** Extracts client connection info from extra handler data */
198209
export function extractClientInfo(extra: ExtraHandlerData): {
199210
address?: string;
200-
port?: number
211+
port?: number;
201212
} {
202213
return {
203-
address: extra?.requestInfo?.remoteAddress ||
204-
extra?.clientAddress ||
205-
extra?.request?.ip ||
206-
extra?.request?.connection?.remoteAddress,
207-
port: extra?.requestInfo?.remotePort ||
208-
extra?.clientPort ||
209-
extra?.request?.connection?.remotePort
214+
address:
215+
extra?.requestInfo?.remoteAddress ||
216+
extra?.clientAddress ||
217+
extra?.request?.ip ||
218+
extra?.request?.connection?.remoteAddress,
219+
port: extra?.requestInfo?.remotePort || extra?.clientPort || extra?.request?.connection?.remotePort,
210220
};
211221
}
212222

@@ -250,70 +260,106 @@ export function buildTypeSpecificAttributes(
250260
return getNotificationAttributes(message.method, params || {});
251261
}
252262

253-
/** Simplified tool result attribute extraction */
254-
export function extractSimpleToolAttributes(result: unknown): Record<string, string | number | boolean> {
255-
const attributes: Record<string, string | number | boolean> = {};
256-
257-
if (typeof result === 'object' && result !== null) {
258-
const resultObj = result as Record<string, unknown>;
259-
260-
// Check if this is an error result
261-
if (typeof resultObj.isError === 'boolean') {
262-
attributes['mcp.tool.result.is_error'] = resultObj.isError;
263-
}
264-
265-
// Extract basic content info
266-
if (Array.isArray(resultObj.content)) {
267-
attributes['mcp.tool.result.content_count'] = resultObj.content.length;
268-
269-
// Extract info from all content items
270-
for (let i = 0; i < resultObj.content.length; i++) {
271-
const item = resultObj.content[i];
272-
if (item && typeof item === 'object' && item !== null) {
273-
const contentItem = item as Record<string, unknown>;
274-
const prefix = resultObj.content.length === 1 ? 'mcp.tool.result' : `mcp.tool.result.${i}`;
275-
276-
// Always capture the content type
277-
if (typeof contentItem.type === 'string') {
278-
attributes[`${prefix}.content_type`] = contentItem.type;
279-
}
280-
281-
// Extract common fields generically
282-
if (typeof contentItem.text === 'string') {
283-
const text = contentItem.text;
284-
attributes[`${prefix}.content`] = text.length > 500 ? `${text.substring(0, 497)}...` : text;
285-
}
286-
287-
if (typeof contentItem.mimeType === 'string') {
288-
attributes[`${prefix}.mime_type`] = contentItem.mimeType;
289-
}
290-
291-
if (typeof contentItem.uri === 'string') {
292-
attributes[`${prefix}.uri`] = contentItem.uri;
293-
}
294-
295-
if (typeof contentItem.name === 'string') {
296-
attributes[`${prefix}.name`] = contentItem.name;
297-
}
298-
299-
if (typeof contentItem.data === 'string') {
300-
attributes[`${prefix}.data_size`] = contentItem.data.length;
301-
}
302-
303-
// For embedded resources, check the nested resource object
304-
if (contentItem.resource && typeof contentItem.resource === 'object') {
305-
const resource = contentItem.resource as Record<string, unknown>;
306-
if (typeof resource.uri === 'string') {
307-
attributes[`${prefix}.resource_uri`] = resource.uri;
308-
}
309-
if (typeof resource.mimeType === 'string') {
310-
attributes[`${prefix}.resource_mime_type`] = resource.mimeType;
311-
}
312-
}
313-
}
263+
/** Get metadata about tool result content array */
264+
function getContentMetadata(content: unknown[]): Record<string, string | number> {
265+
return {
266+
'mcp.tool.result.content_count': content.length,
267+
};
268+
}
269+
270+
/** Build attributes from a single content item */
271+
function buildContentItemAttributes(
272+
contentItem: Record<string, unknown>,
273+
prefix: string,
274+
): Record<string, string | number> {
275+
const attributes: Record<string, string | number> = {};
276+
277+
if (typeof contentItem.type === 'string') {
278+
attributes[`${prefix}.content_type`] = contentItem.type;
279+
}
280+
281+
if (typeof contentItem.text === 'string') {
282+
const text = contentItem.text;
283+
attributes[`${prefix}.content`] = text.length > 500 ? `${text.substring(0, 497)}...` : text;
284+
}
285+
286+
if (typeof contentItem.mimeType === 'string') {
287+
attributes[`${prefix}.mime_type`] = contentItem.mimeType;
288+
}
289+
290+
if (typeof contentItem.uri === 'string') {
291+
attributes[`${prefix}.uri`] = contentItem.uri;
292+
}
293+
294+
if (typeof contentItem.name === 'string') {
295+
attributes[`${prefix}.name`] = contentItem.name;
296+
}
297+
298+
if (typeof contentItem.data === 'string') {
299+
attributes[`${prefix}.data_size`] = contentItem.data.length;
300+
}
301+
302+
return attributes;
303+
}
304+
305+
/** Build attributes from embedded resource object */
306+
function buildEmbeddedResourceAttributes(resource: Record<string, unknown>, prefix: string): Record<string, string> {
307+
const attributes: Record<string, string> = {};
308+
309+
if (typeof resource.uri === 'string') {
310+
attributes[`${prefix}.resource_uri`] = resource.uri;
311+
}
312+
313+
if (typeof resource.mimeType === 'string') {
314+
attributes[`${prefix}.resource_mime_type`] = resource.mimeType;
315+
}
316+
317+
return attributes;
318+
}
319+
320+
/** Build attributes for all content items in the result */
321+
function buildAllContentItemAttributes(content: unknown[]): Record<string, string | number> {
322+
const attributes: Record<string, string | number> = {};
323+
324+
for (let i = 0; i < content.length; i++) {
325+
const item = content[i];
326+
if (item && typeof item === 'object' && item !== null) {
327+
const contentItem = item as Record<string, unknown>;
328+
const prefix = content.length === 1 ? 'mcp.tool.result' : `mcp.tool.result.${i}`;
329+
330+
Object.assign(attributes, buildContentItemAttributes(contentItem, prefix));
331+
332+
if (contentItem.resource && typeof contentItem.resource === 'object') {
333+
const resourceAttrs = buildEmbeddedResourceAttributes(contentItem.resource as Record<string, unknown>, prefix);
334+
Object.assign(attributes, resourceAttrs);
314335
}
315336
}
316337
}
317-
338+
339+
return attributes;
340+
}
341+
342+
/** Extract tool result attributes for span instrumentation */
343+
export function extractToolResultAttributes(result: unknown): Record<string, string | number | boolean> {
344+
let attributes: Record<string, string | number | boolean> = {};
345+
346+
if (typeof result !== 'object' || result === null) {
347+
return attributes;
348+
}
349+
350+
const resultObj = result as Record<string, unknown>;
351+
352+
if (typeof resultObj.isError === 'boolean') {
353+
attributes['mcp.tool.result.is_error'] = resultObj.isError;
354+
}
355+
356+
if (Array.isArray(resultObj.content)) {
357+
attributes = {
358+
...attributes,
359+
...getContentMetadata(resultObj.content),
360+
...buildAllContentItemAttributes(resultObj.content),
361+
};
362+
}
363+
318364
return attributes;
319365
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,4 +117,4 @@ export const MCP_NOTIFICATION_ORIGIN_VALUE = 'auto.mcp.notification';
117117
/**
118118
* Sentry source value for MCP route spans
119119
*/
120-
export const MCP_ROUTE_SOURCE_VALUE = 'route';
120+
export const MCP_ROUTE_SOURCE_VALUE = 'route';

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

Lines changed: 33 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,14 @@ import type { RequestId, SessionId } from './types';
99

1010
// Simplified correlation system that works with or without sessionId
1111
// Maps requestId directly to span data for stateless operation
12-
const requestIdToSpanMap = new Map<RequestId, {
13-
span: Span;
14-
method: string;
15-
startTime: number;
16-
}>();
12+
const requestIdToSpanMap = new Map<
13+
RequestId,
14+
{
15+
span: Span;
16+
method: string;
17+
startTime: number;
18+
}
19+
>();
1720

1821
/**
1922
* Stores span context for later correlation with handler execution
@@ -35,7 +38,7 @@ export function associateContextWithRequestSpan<T>(
3538
): T {
3639
if (extraHandlerData) {
3740
const { requestId } = extraHandlerData;
38-
41+
3942
const spanData = requestIdToSpanMap.get(requestId);
4043
if (!spanData) {
4144
return cb();
@@ -57,32 +60,32 @@ export function completeSpanWithResults(requestId: RequestId, result: unknown):
5760
const spanData = requestIdToSpanMap.get(requestId);
5861
if (spanData) {
5962
const { span, method } = spanData;
60-
61-
const spanWithMethods = span as Span & {
62-
setAttributes: (attrs: Record<string, unknown>) => void;
63+
64+
const spanWithMethods = span as Span & {
65+
setAttributes: (attrs: Record<string, unknown>) => void;
6366
setStatus: (status: { code: number; message: string }) => void;
6467
end: () => void;
6568
};
66-
69+
6770
if (spanWithMethods.setAttributes && method === 'tools/call') {
6871
// Add tool-specific attributes
6972
const toolAttributes = extractToolResultAttributes(result);
7073
spanWithMethods.setAttributes(toolAttributes);
71-
74+
7275
// Set span status based on tool result
7376
if (toolAttributes['mcp.tool.result.is_error']) {
74-
spanWithMethods.setStatus({
77+
spanWithMethods.setStatus({
7578
code: 2, // ERROR
76-
message: 'Tool execution failed'
79+
message: 'Tool execution failed',
7780
});
7881
}
7982
}
80-
83+
8184
// Complete the span
8285
if (spanWithMethods.end) {
8386
spanWithMethods.end();
8487
}
85-
88+
8689
// Clean up correlation
8790
requestIdToSpanMap.delete(requestId);
8891
}
@@ -93,18 +96,21 @@ export function completeSpanWithResults(requestId: RequestId, result: unknown):
9396
*/
9497
export function cleanupAllPendingSpans(): number {
9598
const pendingCount = requestIdToSpanMap.size;
96-
99+
97100
for (const [, spanData] of requestIdToSpanMap) {
98-
const spanWithEnd = spanData.span as Span & { end: () => void; setStatus: (status: { code: number; message: string }) => void };
101+
const spanWithEnd = spanData.span as Span & {
102+
end: () => void;
103+
setStatus: (status: { code: number; message: string }) => void;
104+
};
99105
if (spanWithEnd.setStatus && spanWithEnd.end) {
100-
spanWithEnd.setStatus({
106+
spanWithEnd.setStatus({
101107
code: 2, // ERROR
102-
message: 'Transport closed before request completion'
108+
message: 'Transport closed before request completion',
103109
});
104110
spanWithEnd.end();
105111
}
106112
}
107-
113+
108114
requestIdToSpanMap.clear();
109115
return pendingCount;
110116
}
@@ -114,25 +120,24 @@ export function cleanupAllPendingSpans(): number {
114120
*/
115121
function extractToolResultAttributes(result: unknown): Record<string, string | number | boolean> {
116122
const attributes: Record<string, string | number | boolean> = {};
117-
123+
118124
if (typeof result === 'object' && result !== null) {
119125
const resultObj = result as Record<string, unknown>;
120-
126+
121127
// Check if this is an error result
122128
if (typeof resultObj.isError === 'boolean') {
123129
attributes['mcp.tool.result.is_error'] = resultObj.isError;
124130
}
125-
131+
126132
// Store content as-is (serialized)
127133
if (Array.isArray(resultObj.content)) {
128134
attributes['mcp.tool.result.content_count'] = resultObj.content.length;
129-
135+
130136
const serializedContent = JSON.stringify(resultObj.content);
131-
attributes['mcp.tool.result.content'] = serializedContent.length > 5000
132-
? `${serializedContent.substring(0, 4997)}...`
133-
: serializedContent;
137+
attributes['mcp.tool.result.content'] =
138+
serializedContent.length > 5000 ? `${serializedContent.substring(0, 4997)}...` : serializedContent;
134139
}
135140
}
136-
141+
137142
return attributes;
138143
}

0 commit comments

Comments
 (0)