Skip to content

Commit c959853

Browse files
betegonbetegon
committed
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. (cherry picked from commit 0e05a40)
1 parent 2c36fd3 commit c959853

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)