Skip to content

feat(core): MCP Server - Capture prompt results from prompt function calls #17284

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 13 commits into from
Aug 6, 2025
+526 −279
Merged

feat(core): MCP Server - Capture prompt results from prompt function calls #17284

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
f1fbea0
Capture prompt results in prompt function calls
betegon Aug 1, 2025
df8c23c
Merge branch 'develop' into bete/mcp-server-prompt-results
betegon Aug 1, 2025
507c514
oneliner for ifs
betegon Aug 1, 2025
18d74f1
fix lint
betegon Aug 1, 2025
972d76d
capture ALL messages in result.
betegon Aug 1, 2025
1ae03e5
ref: keep files under 300 code lines
betegon Aug 1, 2025
2c6e5a0
fix lint
betegon Aug 1, 2025
d1e9fdc
avoid inline conditionals
betegon Aug 4, 2025
a0fb99e
use type narrowing instead of casting
betegon Aug 4, 2025
105e8b5
remove length limit
betegon Aug 4, 2025
38a6ea5
avoid spreads
betegon Aug 4, 2025
d2a344f
Merge branch 'develop' into bete/mcp-server-prompt-results
betegon Aug 5, 2025
41556c8
take into account result list attributes as PII
betegon Aug 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
267 changes: 4 additions & 263 deletions packages/core/src/integrations/mcp-server/attributeExtraction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,66 +1,18 @@
/**
* Attribute extraction and building functions for MCP server instrumentation
* Core attribute extraction and building functions for MCP server instrumentation
*/

import { isURLObjectRelative, parseStringToURLObject } from '../../utils/url';
import {
CLIENT_ADDRESS_ATTRIBUTE,
CLIENT_PORT_ATTRIBUTE,
MCP_LOGGING_DATA_TYPE_ATTRIBUTE,
MCP_LOGGING_LEVEL_ATTRIBUTE,
MCP_LOGGING_LOGGER_ATTRIBUTE,
MCP_LOGGING_MESSAGE_ATTRIBUTE,
MCP_PROTOCOL_VERSION_ATTRIBUTE,
MCP_REQUEST_ID_ATTRIBUTE,
MCP_RESOURCE_URI_ATTRIBUTE,
MCP_SERVER_NAME_ATTRIBUTE,
MCP_SERVER_TITLE_ATTRIBUTE,
MCP_SERVER_VERSION_ATTRIBUTE,
MCP_SESSION_ID_ATTRIBUTE,
MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE,
MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE,
MCP_TRANSPORT_ATTRIBUTE,
NETWORK_PROTOCOL_VERSION_ATTRIBUTE,
NETWORK_TRANSPORT_ATTRIBUTE,
} from './attributes';
import { extractTargetInfo, getRequestArguments } from './methodConfig';
import {
getClientInfoForTransport,
getProtocolVersionForTransport,
getSessionDataForTransport,
} from './sessionManagement';
import type {
ExtraHandlerData,
JsonRpcNotification,
JsonRpcRequest,
McpSpanType,
MCPTransport,
PartyInfo,
SessionData,
} from './types';

/**
* Extracts transport types based on transport constructor name
* @param transport - MCP transport instance
* @returns Transport type mapping for span attributes
*/
export function getTransportTypes(transport: MCPTransport): { mcpTransport: string; networkTransport: string } {
const transportName = transport.constructor?.name?.toLowerCase() || '';

if (transportName.includes('stdio')) {
return { mcpTransport: 'stdio', networkTransport: 'pipe' };
}

if (transportName.includes('streamablehttp') || transportName.includes('streamable')) {
return { mcpTransport: 'http', networkTransport: 'tcp' };
}

if (transportName.includes('sse')) {
return { mcpTransport: 'sse', networkTransport: 'tcp' };
}

return { mcpTransport: 'unknown', networkTransport: 'unknown' };
}
import type { JsonRpcNotification, JsonRpcRequest, McpSpanType } from './types';

/**
* Extracts additional attributes for specific notification types
Expand Down Expand Up @@ -138,155 +90,6 @@ export function getNotificationAttributes(
return attributes;
}

/**
* Extracts and validates PartyInfo from an unknown object
* @param obj - Unknown object that might contain party info
* @returns Validated PartyInfo object with only string properties
*/
function extractPartyInfo(obj: unknown): PartyInfo {
const partyInfo: PartyInfo = {};

if (obj && typeof obj === 'object' && obj !== null) {
const source = obj as Record<string, unknown>;
if (typeof source.name === 'string') partyInfo.name = source.name;
if (typeof source.title === 'string') partyInfo.title = source.title;
if (typeof source.version === 'string') partyInfo.version = source.version;
}

return partyInfo;
}

/**
* Extracts session data from "initialize" requests
* @param request - JSON-RPC "initialize" request containing client info and protocol version
* @returns Session data extracted from request parameters including protocol version and client info
*/
export function extractSessionDataFromInitializeRequest(request: JsonRpcRequest): SessionData {
const sessionData: SessionData = {};
if (request.params && typeof request.params === 'object' && request.params !== null) {
const params = request.params as Record<string, unknown>;
if (typeof params.protocolVersion === 'string') {
sessionData.protocolVersion = params.protocolVersion;
}
if (params.clientInfo) {
sessionData.clientInfo = extractPartyInfo(params.clientInfo);
}
}
return sessionData;
}

/**
* Extracts session data from "initialize" response
* @param result - "initialize" response result containing server info and protocol version
* @returns Partial session data extracted from response including protocol version and server info
*/
export function extractSessionDataFromInitializeResponse(result: unknown): Partial<SessionData> {
const sessionData: Partial<SessionData> = {};
if (result && typeof result === 'object') {
const resultObj = result as Record<string, unknown>;
if (typeof resultObj.protocolVersion === 'string') sessionData.protocolVersion = resultObj.protocolVersion;
if (resultObj.serverInfo) {
sessionData.serverInfo = extractPartyInfo(resultObj.serverInfo);
}
}
return sessionData;
}

/**
* Build client attributes from stored client info
* @param transport - MCP transport instance
* @returns Client attributes for span instrumentation
*/
export function getClientAttributes(transport: MCPTransport): Record<string, string> {
const clientInfo = getClientInfoForTransport(transport);
const attributes: Record<string, string> = {};

if (clientInfo?.name) {
attributes['mcp.client.name'] = clientInfo.name;
}
if (clientInfo?.title) {
attributes['mcp.client.title'] = clientInfo.title;
}
if (clientInfo?.version) {
attributes['mcp.client.version'] = clientInfo.version;
}

return attributes;
}

/**
* Build server attributes from stored server info
* @param transport - MCP transport instance
* @returns Server attributes for span instrumentation
*/
export function getServerAttributes(transport: MCPTransport): Record<string, string> {
const serverInfo = getSessionDataForTransport(transport)?.serverInfo;
const attributes: Record<string, string> = {};

if (serverInfo?.name) {
attributes[MCP_SERVER_NAME_ATTRIBUTE] = serverInfo.name;
}
if (serverInfo?.title) {
attributes[MCP_SERVER_TITLE_ATTRIBUTE] = serverInfo.title;
}
if (serverInfo?.version) {
attributes[MCP_SERVER_VERSION_ATTRIBUTE] = serverInfo.version;
}

return attributes;
}

/**
* Extracts client connection info from extra handler data
* @param extra - Extra handler data containing connection info
* @returns Client address and port information
*/
export function extractClientInfo(extra: ExtraHandlerData): {
address?: string;
port?: number;
} {
return {
address:
extra?.requestInfo?.remoteAddress ||
extra?.clientAddress ||
extra?.request?.ip ||
extra?.request?.connection?.remoteAddress,
port: extra?.requestInfo?.remotePort || extra?.clientPort || extra?.request?.connection?.remotePort,
};
}

/**
* Build transport and network attributes
* @param transport - MCP transport instance
* @param extra - Optional extra handler data
* @returns Transport attributes for span instrumentation
*/
export function buildTransportAttributes(
transport: MCPTransport,
extra?: ExtraHandlerData,
): Record<string, string | number> {
const sessionId = transport.sessionId;
const clientInfo = extra ? extractClientInfo(extra) : {};
const { mcpTransport, networkTransport } = getTransportTypes(transport);
const clientAttributes = getClientAttributes(transport);
const serverAttributes = getServerAttributes(transport);
const protocolVersion = getProtocolVersionForTransport(transport);

const attributes = {
...(sessionId && { [MCP_SESSION_ID_ATTRIBUTE]: sessionId }),
...(clientInfo.address && { [CLIENT_ADDRESS_ATTRIBUTE]: clientInfo.address }),
...(clientInfo.port && { [CLIENT_PORT_ATTRIBUTE]: clientInfo.port }),
[MCP_TRANSPORT_ATTRIBUTE]: mcpTransport,
[NETWORK_TRANSPORT_ATTRIBUTE]: networkTransport,
[NETWORK_PROTOCOL_VERSION_ATTRIBUTE]: '2.0',
...(protocolVersion && { [MCP_PROTOCOL_VERSION_ATTRIBUTE]: protocolVersion }),
...clientAttributes,
...serverAttributes,
};

return attributes;
}

/**
* Build type-specific attributes based on message type
* @param type - Span type (request or notification)
Expand All @@ -313,67 +116,5 @@ export function buildTypeSpecificAttributes(
return getNotificationAttributes(message.method, params || {});
}

/**
* Build attributes for tool result content items
* @param content - Array of content items from tool result
* @returns Attributes extracted from each content item including type, text, mime type, URI, and resource info
*/
function buildAllContentItemAttributes(content: unknown[]): Record<string, string | number> {
const attributes: Record<string, string | number> = {
[MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE]: content.length,
};

for (const [i, item] of content.entries()) {
if (typeof item !== 'object' || item === null) continue;

const contentItem = item as Record<string, unknown>;
const prefix = content.length === 1 ? 'mcp.tool.result' : `mcp.tool.result.${i}`;

const safeSet = (key: string, value: unknown): void => {
if (typeof value === 'string') attributes[`${prefix}.${key}`] = value;
};

safeSet('content_type', contentItem.type);
safeSet('mime_type', contentItem.mimeType);
safeSet('uri', contentItem.uri);
safeSet('name', contentItem.name);

if (typeof contentItem.text === 'string') {
const text = contentItem.text;
const maxLength = 500;
attributes[`${prefix}.content`] = text.length > maxLength ? `${text.slice(0, maxLength - 3)}...` : text;
}

if (typeof contentItem.data === 'string') {
attributes[`${prefix}.data_size`] = contentItem.data.length;
}

const resource = contentItem.resource;
if (typeof resource === 'object' && resource !== null) {
const res = resource as Record<string, unknown>;
safeSet('resource_uri', res.uri);
safeSet('resource_mime_type', res.mimeType);
}
}

return attributes;
}

/**
* Extract tool result attributes for span instrumentation
* @param result - Tool execution result
* @returns Attributes extracted from tool result content
*/
export function extractToolResultAttributes(result: unknown): Record<string, string | number | boolean> {
let attributes: Record<string, string | number | boolean> = {};
if (typeof result !== 'object' || result === null) return attributes;

const resultObj = result as Record<string, unknown>;
if (typeof resultObj.isError === 'boolean') {
attributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] = resultObj.isError;
}
if (Array.isArray(resultObj.content)) {
attributes = { ...attributes, ...buildAllContentItemAttributes(resultObj.content) };
}
return attributes;
}
// Re-export buildTransportAttributes for spans.ts
export { buildTransportAttributes } from './sessionExtraction';
22 changes: 22 additions & 0 deletions packages/core/src/integrations/mcp-server/attributes.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,28 @@ export const MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE = 'mcp.tool.result.content_
/** Serialized content of the tool result */
export const MCP_TOOL_RESULT_CONTENT_ATTRIBUTE = 'mcp.tool.result.content';

/** Prefix for tool result attributes that contain sensitive content */
export const MCP_TOOL_RESULT_PREFIX = 'mcp.tool.result';

// =============================================================================
// PROMPT RESULT ATTRIBUTES
// =============================================================================

/** Description of the prompt result */
export const MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE = 'mcp.prompt.result.description';

/** Number of messages in the prompt result */
export const MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE = 'mcp.prompt.result.message_count';

/** Role of the message in the prompt result (for single message results) */
export const MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE = 'mcp.prompt.result.message_role';

/** Content of the message in the prompt result (for single message results) */
export const MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE = 'mcp.prompt.result.message_content';

/** Prefix for prompt result attributes that contain sensitive content */
export const MCP_PROMPT_RESULT_PREFIX = 'mcp.prompt.result';

// =============================================================================
// REQUEST ARGUMENT ATTRIBUTES
// =============================================================================
Expand Down
13 changes: 11 additions & 2 deletions packages/core/src/integrations/mcp-server/correlation.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@
import { getClient } from '../../currentScopes';
import { SPAN_STATUS_ERROR } from '../../tracing';
import type { Span } from '../../types-hoist/span';
import { extractToolResultAttributes } from './attributeExtraction';
import { filterMcpPiiFromSpanData } from './piiFiltering';
import { extractPromptResultAttributes, extractToolResultAttributes } from './resultExtraction';
import type { MCPTransport, RequestId, RequestSpanMapValue } from './types';

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

span.setAttributes(toolAttributes);
} else if (method === 'prompts/get') {
const rawPromptAttributes = extractPromptResultAttributes(result);
const client = getClient();
const sendDefaultPii = Boolean(client?.getOptions().sendDefaultPii);
const promptAttributes = filterMcpPiiFromSpanData(rawPromptAttributes, sendDefaultPii);

span.setAttributes(promptAttributes);
}

span.end();
Expand All @@ -83,7 +90,9 @@ export function completeSpanWithResults(transport: MCPTransport, requestId: Requ
*/
export function cleanupPendingSpansForTransport(transport: MCPTransport): number {
const spanMap = transportToSpanMap.get(transport);
if (!spanMap) return 0;
if (!spanMap) {
return 0;
}

const pendingCount = spanMap.size;

Expand Down
Loading
Loading