Skip to content

Commit 0e05a40

Browse files
betegonbetegon
authored
feat(core): MCP Server - Capture prompt results from prompt function calls (#17284)
closes #17283 includes these attributes for `mcp.server` spans: - `mcp.prompt.result.description` - `mcp.prompt.result.message_content` - `mcp.prompt.result.message_role` - `mcp.prompt.result.message_count` Example: <img width="835" height="300" alt="Screenshot 2025-08-01 at 12 40 46" src="https://github.com/user-attachments/assets/592d876b-807a-4f3e-a9b2-406e10f5a83d" /> Needed to make `attributeExtraction.ts` <300 lines of code (requirement) so it's now split between `sessionExtraction.ts`, `sessionExtraction.ts` and `resultExtraction.ts`. So changes explained so it's easier to review: - The only function this PR adds is `extractPromptResultAttributes` inside `resultExtraction.ts`. - It adds the prompt results as PII in `piiFiltering.ts`. Just add them to the `set`. - adds a `else if (method === 'prompts/get')` to execute the `extractPromptResultAttributes` function. - adds a test that checks we're capturing the results and updates the PII test to check PII result attributes are being removed if sending PII is not enabled.
1 parent ed07836 commit 0e05a40

File tree

12 files changed

+526
-279
lines changed

12 files changed

+526
-279
lines changed

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

Lines changed: 4 additions & 263 deletions
Original file line numberDiff line numberDiff line change
@@ -1,66 +1,18 @@
11
/**
2-
* Attribute extraction and building functions for MCP server instrumentation
2+
* Core attribute extraction and building functions for MCP server instrumentation
33
*/
44

55
import { isURLObjectRelative, parseStringToURLObject } from '../../utils/url';
66
import {
7-
CLIENT_ADDRESS_ATTRIBUTE,
8-
CLIENT_PORT_ATTRIBUTE,
97
MCP_LOGGING_DATA_TYPE_ATTRIBUTE,
108
MCP_LOGGING_LEVEL_ATTRIBUTE,
119
MCP_LOGGING_LOGGER_ATTRIBUTE,
1210
MCP_LOGGING_MESSAGE_ATTRIBUTE,
13-
MCP_PROTOCOL_VERSION_ATTRIBUTE,
1411
MCP_REQUEST_ID_ATTRIBUTE,
1512
MCP_RESOURCE_URI_ATTRIBUTE,
16-
MCP_SERVER_NAME_ATTRIBUTE,
17-
MCP_SERVER_TITLE_ATTRIBUTE,
18-
MCP_SERVER_VERSION_ATTRIBUTE,
19-
MCP_SESSION_ID_ATTRIBUTE,
20-
MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE,
21-
MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE,
22-
MCP_TRANSPORT_ATTRIBUTE,
23-
NETWORK_PROTOCOL_VERSION_ATTRIBUTE,
24-
NETWORK_TRANSPORT_ATTRIBUTE,
2513
} from './attributes';
2614
import { extractTargetInfo, getRequestArguments } from './methodConfig';
27-
import {
28-
getClientInfoForTransport,
29-
getProtocolVersionForTransport,
30-
getSessionDataForTransport,
31-
} from './sessionManagement';
32-
import type {
33-
ExtraHandlerData,
34-
JsonRpcNotification,
35-
JsonRpcRequest,
36-
McpSpanType,
37-
MCPTransport,
38-
PartyInfo,
39-
SessionData,
40-
} from './types';
41-
42-
/**
43-
* Extracts transport types based on transport constructor name
44-
* @param transport - MCP transport instance
45-
* @returns Transport type mapping for span attributes
46-
*/
47-
export function getTransportTypes(transport: MCPTransport): { mcpTransport: string; networkTransport: string } {
48-
const transportName = transport.constructor?.name?.toLowerCase() || '';
49-
50-
if (transportName.includes('stdio')) {
51-
return { mcpTransport: 'stdio', networkTransport: 'pipe' };
52-
}
53-
54-
if (transportName.includes('streamablehttp') || transportName.includes('streamable')) {
55-
return { mcpTransport: 'http', networkTransport: 'tcp' };
56-
}
57-
58-
if (transportName.includes('sse')) {
59-
return { mcpTransport: 'sse', networkTransport: 'tcp' };
60-
}
61-
62-
return { mcpTransport: 'unknown', networkTransport: 'unknown' };
63-
}
15+
import type { JsonRpcNotification, JsonRpcRequest, McpSpanType } from './types';
6416

6517
/**
6618
* Extracts additional attributes for specific notification types
@@ -138,155 +90,6 @@ export function getNotificationAttributes(
13890
return attributes;
13991
}
14092

141-
/**
142-
* Extracts and validates PartyInfo from an unknown object
143-
* @param obj - Unknown object that might contain party info
144-
* @returns Validated PartyInfo object with only string properties
145-
*/
146-
function extractPartyInfo(obj: unknown): PartyInfo {
147-
const partyInfo: PartyInfo = {};
148-
149-
if (obj && typeof obj === 'object' && obj !== null) {
150-
const source = obj as Record<string, unknown>;
151-
if (typeof source.name === 'string') partyInfo.name = source.name;
152-
if (typeof source.title === 'string') partyInfo.title = source.title;
153-
if (typeof source.version === 'string') partyInfo.version = source.version;
154-
}
155-
156-
return partyInfo;
157-
}
158-
159-
/**
160-
* Extracts session data from "initialize" requests
161-
* @param request - JSON-RPC "initialize" request containing client info and protocol version
162-
* @returns Session data extracted from request parameters including protocol version and client info
163-
*/
164-
export function extractSessionDataFromInitializeRequest(request: JsonRpcRequest): SessionData {
165-
const sessionData: SessionData = {};
166-
if (request.params && typeof request.params === 'object' && request.params !== null) {
167-
const params = request.params as Record<string, unknown>;
168-
if (typeof params.protocolVersion === 'string') {
169-
sessionData.protocolVersion = params.protocolVersion;
170-
}
171-
if (params.clientInfo) {
172-
sessionData.clientInfo = extractPartyInfo(params.clientInfo);
173-
}
174-
}
175-
return sessionData;
176-
}
177-
178-
/**
179-
* Extracts session data from "initialize" response
180-
* @param result - "initialize" response result containing server info and protocol version
181-
* @returns Partial session data extracted from response including protocol version and server info
182-
*/
183-
export function extractSessionDataFromInitializeResponse(result: unknown): Partial<SessionData> {
184-
const sessionData: Partial<SessionData> = {};
185-
if (result && typeof result === 'object') {
186-
const resultObj = result as Record<string, unknown>;
187-
if (typeof resultObj.protocolVersion === 'string') sessionData.protocolVersion = resultObj.protocolVersion;
188-
if (resultObj.serverInfo) {
189-
sessionData.serverInfo = extractPartyInfo(resultObj.serverInfo);
190-
}
191-
}
192-
return sessionData;
193-
}
194-
195-
/**
196-
* Build client attributes from stored client info
197-
* @param transport - MCP transport instance
198-
* @returns Client attributes for span instrumentation
199-
*/
200-
export function getClientAttributes(transport: MCPTransport): Record<string, string> {
201-
const clientInfo = getClientInfoForTransport(transport);
202-
const attributes: Record<string, string> = {};
203-
204-
if (clientInfo?.name) {
205-
attributes['mcp.client.name'] = clientInfo.name;
206-
}
207-
if (clientInfo?.title) {
208-
attributes['mcp.client.title'] = clientInfo.title;
209-
}
210-
if (clientInfo?.version) {
211-
attributes['mcp.client.version'] = clientInfo.version;
212-
}
213-
214-
return attributes;
215-
}
216-
217-
/**
218-
* Build server attributes from stored server info
219-
* @param transport - MCP transport instance
220-
* @returns Server attributes for span instrumentation
221-
*/
222-
export function getServerAttributes(transport: MCPTransport): Record<string, string> {
223-
const serverInfo = getSessionDataForTransport(transport)?.serverInfo;
224-
const attributes: Record<string, string> = {};
225-
226-
if (serverInfo?.name) {
227-
attributes[MCP_SERVER_NAME_ATTRIBUTE] = serverInfo.name;
228-
}
229-
if (serverInfo?.title) {
230-
attributes[MCP_SERVER_TITLE_ATTRIBUTE] = serverInfo.title;
231-
}
232-
if (serverInfo?.version) {
233-
attributes[MCP_SERVER_VERSION_ATTRIBUTE] = serverInfo.version;
234-
}
235-
236-
return attributes;
237-
}
238-
239-
/**
240-
* Extracts client connection info from extra handler data
241-
* @param extra - Extra handler data containing connection info
242-
* @returns Client address and port information
243-
*/
244-
export function extractClientInfo(extra: ExtraHandlerData): {
245-
address?: string;
246-
port?: number;
247-
} {
248-
return {
249-
address:
250-
extra?.requestInfo?.remoteAddress ||
251-
extra?.clientAddress ||
252-
extra?.request?.ip ||
253-
extra?.request?.connection?.remoteAddress,
254-
port: extra?.requestInfo?.remotePort || extra?.clientPort || extra?.request?.connection?.remotePort,
255-
};
256-
}
257-
258-
/**
259-
* Build transport and network attributes
260-
* @param transport - MCP transport instance
261-
* @param extra - Optional extra handler data
262-
* @returns Transport attributes for span instrumentation
263-
*/
264-
export function buildTransportAttributes(
265-
transport: MCPTransport,
266-
extra?: ExtraHandlerData,
267-
): Record<string, string | number> {
268-
const sessionId = transport.sessionId;
269-
const clientInfo = extra ? extractClientInfo(extra) : {};
270-
const { mcpTransport, networkTransport } = getTransportTypes(transport);
271-
const clientAttributes = getClientAttributes(transport);
272-
const serverAttributes = getServerAttributes(transport);
273-
const protocolVersion = getProtocolVersionForTransport(transport);
274-
275-
const attributes = {
276-
...(sessionId && { [MCP_SESSION_ID_ATTRIBUTE]: sessionId }),
277-
...(clientInfo.address && { [CLIENT_ADDRESS_ATTRIBUTE]: clientInfo.address }),
278-
...(clientInfo.port && { [CLIENT_PORT_ATTRIBUTE]: clientInfo.port }),
279-
[MCP_TRANSPORT_ATTRIBUTE]: mcpTransport,
280-
[NETWORK_TRANSPORT_ATTRIBUTE]: networkTransport,
281-
[NETWORK_PROTOCOL_VERSION_ATTRIBUTE]: '2.0',
282-
...(protocolVersion && { [MCP_PROTOCOL_VERSION_ATTRIBUTE]: protocolVersion }),
283-
...clientAttributes,
284-
...serverAttributes,
285-
};
286-
287-
return attributes;
288-
}
289-
29093
/**
29194
* Build type-specific attributes based on message type
29295
* @param type - Span type (request or notification)
@@ -313,67 +116,5 @@ export function buildTypeSpecificAttributes(
313116
return getNotificationAttributes(message.method, params || {});
314117
}
315118

316-
/**
317-
* Build attributes for tool result content items
318-
* @param content - Array of content items from tool result
319-
* @returns Attributes extracted from each content item including type, text, mime type, URI, and resource info
320-
*/
321-
function buildAllContentItemAttributes(content: unknown[]): Record<string, string | number> {
322-
const attributes: Record<string, string | number> = {
323-
[MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE]: content.length,
324-
};
325-
326-
for (const [i, item] of content.entries()) {
327-
if (typeof item !== 'object' || item === null) continue;
328-
329-
const contentItem = item as Record<string, unknown>;
330-
const prefix = content.length === 1 ? 'mcp.tool.result' : `mcp.tool.result.${i}`;
331-
332-
const safeSet = (key: string, value: unknown): void => {
333-
if (typeof value === 'string') attributes[`${prefix}.${key}`] = value;
334-
};
335-
336-
safeSet('content_type', contentItem.type);
337-
safeSet('mime_type', contentItem.mimeType);
338-
safeSet('uri', contentItem.uri);
339-
safeSet('name', contentItem.name);
340-
341-
if (typeof contentItem.text === 'string') {
342-
const text = contentItem.text;
343-
const maxLength = 500;
344-
attributes[`${prefix}.content`] = text.length > maxLength ? `${text.slice(0, maxLength - 3)}...` : text;
345-
}
346-
347-
if (typeof contentItem.data === 'string') {
348-
attributes[`${prefix}.data_size`] = contentItem.data.length;
349-
}
350-
351-
const resource = contentItem.resource;
352-
if (typeof resource === 'object' && resource !== null) {
353-
const res = resource as Record<string, unknown>;
354-
safeSet('resource_uri', res.uri);
355-
safeSet('resource_mime_type', res.mimeType);
356-
}
357-
}
358-
359-
return attributes;
360-
}
361-
362-
/**
363-
* Extract tool result attributes for span instrumentation
364-
* @param result - Tool execution result
365-
* @returns Attributes extracted from tool result content
366-
*/
367-
export function extractToolResultAttributes(result: unknown): Record<string, string | number | boolean> {
368-
let attributes: Record<string, string | number | boolean> = {};
369-
if (typeof result !== 'object' || result === null) return attributes;
370-
371-
const resultObj = result as Record<string, unknown>;
372-
if (typeof resultObj.isError === 'boolean') {
373-
attributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] = resultObj.isError;
374-
}
375-
if (Array.isArray(resultObj.content)) {
376-
attributes = { ...attributes, ...buildAllContentItemAttributes(resultObj.content) };
377-
}
378-
return attributes;
379-
}
119+
// Re-export buildTransportAttributes for spans.ts
120+
export { buildTransportAttributes } from './sessionExtraction';

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,28 @@ export const MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE = 'mcp.tool.result.content_
7676
/** Serialized content of the tool result */
7777
export const MCP_TOOL_RESULT_CONTENT_ATTRIBUTE = 'mcp.tool.result.content';
7878

79+
/** Prefix for tool result attributes that contain sensitive content */
80+
export const MCP_TOOL_RESULT_PREFIX = 'mcp.tool.result';
81+
82+
// =============================================================================
83+
// PROMPT RESULT ATTRIBUTES
84+
// =============================================================================
85+
86+
/** Description of the prompt result */
87+
export const MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE = 'mcp.prompt.result.description';
88+
89+
/** Number of messages in the prompt result */
90+
export const MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE = 'mcp.prompt.result.message_count';
91+
92+
/** Role of the message in the prompt result (for single message results) */
93+
export const MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE = 'mcp.prompt.result.message_role';
94+
95+
/** Content of the message in the prompt result (for single message results) */
96+
export const MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE = 'mcp.prompt.result.message_content';
97+
98+
/** Prefix for prompt result attributes that contain sensitive content */
99+
export const MCP_PROMPT_RESULT_PREFIX = 'mcp.prompt.result';
100+
79101
// =============================================================================
80102
// REQUEST ARGUMENT ATTRIBUTES
81103
// =============================================================================

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99
import { getClient } from '../../currentScopes';
1010
import { SPAN_STATUS_ERROR } from '../../tracing';
1111
import type { Span } from '../../types-hoist/span';
12-
import { extractToolResultAttributes } from './attributeExtraction';
1312
import { filterMcpPiiFromSpanData } from './piiFiltering';
13+
import { extractPromptResultAttributes, extractToolResultAttributes } from './resultExtraction';
1414
import type { MCPTransport, RequestId, RequestSpanMapValue } from './types';
1515

1616
/**
@@ -69,6 +69,13 @@ export function completeSpanWithResults(transport: MCPTransport, requestId: Requ
6969
const toolAttributes = filterMcpPiiFromSpanData(rawToolAttributes, sendDefaultPii);
7070

7171
span.setAttributes(toolAttributes);
72+
} else if (method === 'prompts/get') {
73+
const rawPromptAttributes = extractPromptResultAttributes(result);
74+
const client = getClient();
75+
const sendDefaultPii = Boolean(client?.getOptions().sendDefaultPii);
76+
const promptAttributes = filterMcpPiiFromSpanData(rawPromptAttributes, sendDefaultPii);
77+
78+
span.setAttributes(promptAttributes);
7279
}
7380

7481
span.end();
@@ -83,7 +90,9 @@ export function completeSpanWithResults(transport: MCPTransport, requestId: Requ
8390
*/
8491
export function cleanupPendingSpansForTransport(transport: MCPTransport): number {
8592
const spanMap = transportToSpanMap.get(transport);
86-
if (!spanMap) return 0;
93+
if (!spanMap) {
94+
return 0;
95+
}
8796

8897
const pendingCount = spanMap.size;
8998

0 commit comments

Comments
 (0)