From f1fbea0a518eddccea570a5007f4c4133f94e7a5 Mon Sep 17 00:00:00 2001 From: betegon Date: Fri, 1 Aug 2025 12:02:53 +0200 Subject: [PATCH 01/11] Capture prompt results in prompt function calls --- .../mcp-server/attributeExtraction.ts | 44 +++++++++++ .../src/integrations/mcp-server/attributes.ts | 16 ++++ .../integrations/mcp-server/correlation.ts | 9 ++- .../integrations/mcp-server/piiFiltering.ts | 4 + .../mcp-server/piiFiltering.test.ts | 8 +- .../mcp-server/semanticConventions.test.ts | 74 +++++++++++++++++++ 6 files changed, 153 insertions(+), 2 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/attributeExtraction.ts b/packages/core/src/integrations/mcp-server/attributeExtraction.ts index 68eade987a08..68382b3e4f62 100644 --- a/packages/core/src/integrations/mcp-server/attributeExtraction.ts +++ b/packages/core/src/integrations/mcp-server/attributeExtraction.ts @@ -10,6 +10,10 @@ import { MCP_LOGGING_LEVEL_ATTRIBUTE, MCP_LOGGING_LOGGER_ATTRIBUTE, MCP_LOGGING_MESSAGE_ATTRIBUTE, + MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE, MCP_PROTOCOL_VERSION_ATTRIBUTE, MCP_REQUEST_ID_ATTRIBUTE, MCP_RESOURCE_URI_ATTRIBUTE, @@ -377,3 +381,43 @@ export function extractToolResultAttributes(result: unknown): Record { + const attributes: Record = {}; + if (typeof result !== 'object' || result === null) return attributes; + + const resultObj = result as Record; + + if (typeof resultObj.description === 'string') { + attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; + } + + if (Array.isArray(resultObj.messages)) { + attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length; + + if (resultObj.messages.length > 0) { + const message = resultObj.messages[0]; + if (typeof message === 'object' && message !== null) { + const messageObj = message as Record; + + if (typeof messageObj.role === 'string') { + attributes[MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE] = messageObj.role; + } + + if (typeof messageObj.content === 'object' && messageObj.content !== null) { + const content = messageObj.content as Record; + if (typeof content.text === 'string') { + attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text; + } + } + } + } + } + + return attributes; +} diff --git a/packages/core/src/integrations/mcp-server/attributes.ts b/packages/core/src/integrations/mcp-server/attributes.ts index 074bd09b7bdf..e4518d63a9a0 100644 --- a/packages/core/src/integrations/mcp-server/attributes.ts +++ b/packages/core/src/integrations/mcp-server/attributes.ts @@ -76,6 +76,22 @@ 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'; +// ============================================================================= +// 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 prompt message (for single message prompts) */ +export const MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE = 'mcp.prompt.result.message_role'; + +/** Content of the prompt message (for single message prompts) */ +export const MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE = 'mcp.prompt.result.message_content'; + // ============================================================================= // REQUEST ARGUMENT ATTRIBUTES // ============================================================================= diff --git a/packages/core/src/integrations/mcp-server/correlation.ts b/packages/core/src/integrations/mcp-server/correlation.ts index 7f00341bdd5a..93f9c96280da 100644 --- a/packages/core/src/integrations/mcp-server/correlation.ts +++ b/packages/core/src/integrations/mcp-server/correlation.ts @@ -9,7 +9,7 @@ import { getClient } from '../../currentScopes'; import { SPAN_STATUS_ERROR } from '../../tracing'; import type { Span } from '../../types-hoist/span'; -import { extractToolResultAttributes } from './attributeExtraction'; +import { extractPromptResultAttributes, extractToolResultAttributes } from './attributeExtraction'; import { filterMcpPiiFromSpanData } from './piiFiltering'; import type { MCPTransport, RequestId, RequestSpanMapValue } from './types'; @@ -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(); diff --git a/packages/core/src/integrations/mcp-server/piiFiltering.ts b/packages/core/src/integrations/mcp-server/piiFiltering.ts index 654427ca2d6d..0161b926cb23 100644 --- a/packages/core/src/integrations/mcp-server/piiFiltering.ts +++ b/packages/core/src/integrations/mcp-server/piiFiltering.ts @@ -9,6 +9,8 @@ import { CLIENT_ADDRESS_ATTRIBUTE, CLIENT_PORT_ATTRIBUTE, MCP_LOGGING_MESSAGE_ATTRIBUTE, + MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE, MCP_REQUEST_ARGUMENT, MCP_RESOURCE_URI_ATTRIBUTE, MCP_TOOL_RESULT_CONTENT_ATTRIBUTE, @@ -22,6 +24,8 @@ const PII_ATTRIBUTES = new Set([ CLIENT_ADDRESS_ATTRIBUTE, CLIENT_PORT_ATTRIBUTE, MCP_LOGGING_MESSAGE_ATTRIBUTE, + MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE, MCP_RESOURCE_URI_ATTRIBUTE, MCP_TOOL_RESULT_CONTENT_ATTRIBUTE, ]); diff --git a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts index 14f803b28ccc..6193019fc0ff 100644 --- a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts @@ -124,7 +124,7 @@ describe('MCP Server PII Filtering', () => { setAttributes: vi.fn(), setStatus: vi.fn(), end: vi.fn(), - } as any; + } as unknown as ReturnType; startInactiveSpanSpy.mockReturnValueOnce(mockSpan); const toolCallRequest = { @@ -163,6 +163,8 @@ describe('MCP Server PII Filtering', () => { 'client.port': 54321, 'mcp.request.argument.location': '"San Francisco"', 'mcp.tool.result.content': 'Weather data: 18°C', + 'mcp.prompt.result.description': 'Code review prompt for sensitive analysis', + 'mcp.prompt.result.message_content': 'Please review this confidential code.', 'mcp.logging.message': 'User requested weather', 'mcp.resource.uri': 'file:///private/docs/secret.txt', 'mcp.method.name': 'tools/call', // Non-PII should remain @@ -180,6 +182,8 @@ describe('MCP Server PII Filtering', () => { 'mcp.request.argument.location': '"San Francisco"', 'mcp.request.argument.units': '"celsius"', 'mcp.tool.result.content': 'Weather data: 18°C', + 'mcp.prompt.result.description': 'Code review prompt for sensitive analysis', + 'mcp.prompt.result.message_content': 'Please review this confidential code.', 'mcp.logging.message': 'User requested weather', 'mcp.resource.uri': 'file:///private/docs/secret.txt', 'mcp.method.name': 'tools/call', // Non-PII should remain @@ -193,6 +197,8 @@ describe('MCP Server PII Filtering', () => { expect(result).not.toHaveProperty('mcp.request.argument.location'); expect(result).not.toHaveProperty('mcp.request.argument.units'); expect(result).not.toHaveProperty('mcp.tool.result.content'); + expect(result).not.toHaveProperty('mcp.prompt.result.description'); + expect(result).not.toHaveProperty('mcp.prompt.result.message_content'); expect(result).not.toHaveProperty('mcp.logging.message'); expect(result).not.toHaveProperty('mcp.resource.uri'); diff --git a/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts b/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts index 7b110a0b2756..a5d802f8aa4f 100644 --- a/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts @@ -434,5 +434,79 @@ describe('MCP Server Semantic Conventions', () => { expect(setStatusSpy).not.toHaveBeenCalled(); expect(endSpy).toHaveBeenCalled(); }); + + it('should instrument prompt call results and complete span with enriched attributes', async () => { + await wrappedMcpServer.connect(mockTransport); + + const setAttributesSpy = vi.fn(); + const setStatusSpy = vi.fn(); + const endSpy = vi.fn(); + const mockSpan = { + setAttributes: setAttributesSpy, + setStatus: setStatusSpy, + end: endSpy, + }; + startInactiveSpanSpy.mockReturnValueOnce( + mockSpan as unknown as ReturnType, + ); + + const promptCallRequest = { + jsonrpc: '2.0', + method: 'prompts/get', + id: 'req-prompt-result', + params: { + name: 'code-review', + arguments: { language: 'typescript', complexity: 'high' }, + }, + }; + + mockTransport.onmessage?.(promptCallRequest, {}); + + expect(startInactiveSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'prompts/get code-review', + op: 'mcp.server', + forceTransaction: true, + attributes: expect.objectContaining({ + 'mcp.method.name': 'prompts/get', + 'mcp.prompt.name': 'code-review', + 'mcp.request.id': 'req-prompt-result', + }), + }), + ); + + const promptResponse = { + jsonrpc: '2.0', + id: 'req-prompt-result', + result: { + description: 'Code review prompt for TypeScript with high complexity analysis', + messages: [ + { + role: 'user', + content: { + type: 'text', + text: 'Please review this TypeScript code for complexity and best practices.', + }, + }, + ], + }, + }; + + mockTransport.send?.(promptResponse); + + expect(setAttributesSpy).toHaveBeenCalledWith( + expect.objectContaining({ + 'mcp.prompt.result.description': 'Code review prompt for TypeScript with high complexity analysis', + 'mcp.prompt.result.message_count': 1, + 'mcp.prompt.result.message_role': 'user', + 'mcp.prompt.result.message_content': 'Please review this TypeScript code for complexity and best practices.', + }), + ); + + expect(setStatusSpy).not.toHaveBeenCalled(); + expect(endSpy).toHaveBeenCalled(); + }); + + }); }); From 507c514ab12cd354d81ef70659d98eaf3708ac6c Mon Sep 17 00:00:00 2001 From: betegon Date: Fri, 1 Aug 2025 12:14:53 +0200 Subject: [PATCH 02/11] oneliner for ifs --- .../integrations/mcp-server/attributeExtraction.ts | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/attributeExtraction.ts b/packages/core/src/integrations/mcp-server/attributeExtraction.ts index 68382b3e4f62..63c8171a19e8 100644 --- a/packages/core/src/integrations/mcp-server/attributeExtraction.ts +++ b/packages/core/src/integrations/mcp-server/attributeExtraction.ts @@ -393,9 +393,7 @@ export function extractPromptResultAttributes(result: unknown): Record; - if (typeof resultObj.description === 'string') { - attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; - } + if (typeof resultObj.description === 'string') attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; if (Array.isArray(resultObj.messages)) { attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length; @@ -405,15 +403,11 @@ export function extractPromptResultAttributes(result: unknown): Record; - if (typeof messageObj.role === 'string') { - attributes[MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE] = messageObj.role; - } + if (typeof messageObj.role === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE] = messageObj.role; if (typeof messageObj.content === 'object' && messageObj.content !== null) { const content = messageObj.content as Record; - if (typeof content.text === 'string') { - attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text; - } + if (typeof content.text === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text; } } } From 18d74f1cfdf3c2b7a4190e32aa3e20339596eeac Mon Sep 17 00:00:00 2001 From: betegon Date: Fri, 1 Aug 2025 12:50:34 +0200 Subject: [PATCH 03/11] fix lint --- .../mcp-server/attributeExtraction.ts | 15 ++++++++------- .../mcp-server/semanticConventions.test.ts | 2 -- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/attributeExtraction.ts b/packages/core/src/integrations/mcp-server/attributeExtraction.ts index 63c8171a19e8..c5f62b3e25a9 100644 --- a/packages/core/src/integrations/mcp-server/attributeExtraction.ts +++ b/packages/core/src/integrations/mcp-server/attributeExtraction.ts @@ -392,19 +392,20 @@ export function extractPromptResultAttributes(result: unknown): Record; - - if (typeof resultObj.description === 'string') attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; - + + if (typeof resultObj.description === 'string') + attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; + if (Array.isArray(resultObj.messages)) { attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length; - + if (resultObj.messages.length > 0) { const message = resultObj.messages[0]; if (typeof message === 'object' && message !== null) { const messageObj = message as Record; - + if (typeof messageObj.role === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE] = messageObj.role; - + if (typeof messageObj.content === 'object' && messageObj.content !== null) { const content = messageObj.content as Record; if (typeof content.text === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text; @@ -412,6 +413,6 @@ export function extractPromptResultAttributes(result: unknown): Record { expect(setStatusSpy).not.toHaveBeenCalled(); expect(endSpy).toHaveBeenCalled(); }); - - }); }); From 972d76dad3bd1a07bfb981a8445ed87be2eac2bc Mon Sep 17 00:00:00 2001 From: betegon Date: Fri, 1 Aug 2025 14:26:50 +0200 Subject: [PATCH 04/11] capture ALL messages in result. --- .../mcp-server/attributeExtraction.ts | 30 ++++++++++------ .../src/integrations/mcp-server/attributes.ts | 4 +-- .../mcp-server/piiFiltering.test.ts | 35 +++++++++++++++++-- 3 files changed, 54 insertions(+), 15 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/attributeExtraction.ts b/packages/core/src/integrations/mcp-server/attributeExtraction.ts index c5f62b3e25a9..11b220901033 100644 --- a/packages/core/src/integrations/mcp-server/attributeExtraction.ts +++ b/packages/core/src/integrations/mcp-server/attributeExtraction.ts @@ -11,9 +11,7 @@ import { MCP_LOGGING_LOGGER_ATTRIBUTE, MCP_LOGGING_MESSAGE_ATTRIBUTE, MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, - MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE, MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE, - MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE, MCP_PROTOCOL_VERSION_ATTRIBUTE, MCP_REQUEST_ID_ATTRIBUTE, MCP_RESOURCE_URI_ATTRIBUTE, @@ -399,16 +397,28 @@ export function extractPromptResultAttributes(result: unknown): Record 0) { - const message = resultObj.messages[0]; - if (typeof message === 'object' && message !== null) { - const messageObj = message as Record; + // Extract attributes for each message + const messages = resultObj.messages; + for (const [i, message] of messages.entries()) { + if (typeof message !== 'object' || message === null) continue; - if (typeof messageObj.role === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE] = messageObj.role; + const messageObj = message as Record; + const prefix = messages.length === 1 ? 'mcp.prompt.result' : `mcp.prompt.result.${i}`; - if (typeof messageObj.content === 'object' && messageObj.content !== null) { - const content = messageObj.content as Record; - if (typeof content.text === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text; + const safeSet = (key: string, value: unknown): void => { + if (typeof value === 'string') { + const attrName = messages.length === 1 ? `${prefix}.message_${key}` : `${prefix}.${key}`; + attributes[attrName] = value; + } + }; + + safeSet('role', messageObj.role); + + if (typeof messageObj.content === 'object' && messageObj.content !== null) { + const content = messageObj.content as Record; + if (typeof content.text === 'string') { + const attrName = messages.length === 1 ? `${prefix}.message_content` : `${prefix}.content`; + attributes[attrName] = content.text; } } } diff --git a/packages/core/src/integrations/mcp-server/attributes.ts b/packages/core/src/integrations/mcp-server/attributes.ts index e4518d63a9a0..faad50052347 100644 --- a/packages/core/src/integrations/mcp-server/attributes.ts +++ b/packages/core/src/integrations/mcp-server/attributes.ts @@ -86,10 +86,10 @@ export const MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE = 'mcp.prompt.result.descri /** Number of messages in the prompt result */ export const MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE = 'mcp.prompt.result.message_count'; -/** Role of the prompt message (for single message prompts) */ +/** 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 prompt message (for single message prompts) */ +/** 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'; // ============================================================================= diff --git a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts index 6193019fc0ff..e79220f7530d 100644 --- a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts @@ -147,7 +147,7 @@ describe('MCP Server PII Filtering', () => { mockTransport.send?.(toolResponse); - // Tool result content should be filtered out + // Tool result content should be filtered out, but metadata should remain const setAttributesCall = mockSpan.setAttributes.mock.calls[0]?.[0]; expect(setAttributesCall).toBeDefined(); expect(setAttributesCall).not.toHaveProperty('mcp.tool.result.content'); @@ -163,8 +163,11 @@ describe('MCP Server PII Filtering', () => { 'client.port': 54321, 'mcp.request.argument.location': '"San Francisco"', 'mcp.tool.result.content': 'Weather data: 18°C', + 'mcp.tool.result.content_count': 1, 'mcp.prompt.result.description': 'Code review prompt for sensitive analysis', 'mcp.prompt.result.message_content': 'Please review this confidential code.', + 'mcp.prompt.result.message_count': 1, + 'mcp.resource.result.content': 'Sensitive resource content', 'mcp.logging.message': 'User requested weather', 'mcp.resource.uri': 'file:///private/docs/secret.txt', 'mcp.method.name': 'tools/call', // Non-PII should remain @@ -182,8 +185,16 @@ describe('MCP Server PII Filtering', () => { 'mcp.request.argument.location': '"San Francisco"', 'mcp.request.argument.units': '"celsius"', 'mcp.tool.result.content': 'Weather data: 18°C', + 'mcp.tool.result.content_count': 1, 'mcp.prompt.result.description': 'Code review prompt for sensitive analysis', - 'mcp.prompt.result.message_content': 'Please review this confidential code.', + 'mcp.prompt.result.message_count': 2, + 'mcp.prompt.result.0.role': 'user', + 'mcp.prompt.result.0.content': 'Sensitive prompt content', + 'mcp.prompt.result.1.role': 'assistant', + 'mcp.prompt.result.1.content': 'Another sensitive response', + 'mcp.resource.result.content_count': 1, + 'mcp.resource.result.uri': 'file:///private/file.txt', + 'mcp.resource.result.content': 'Sensitive resource content', 'mcp.logging.message': 'User requested weather', 'mcp.resource.uri': 'file:///private/docs/secret.txt', 'mcp.method.name': 'tools/call', // Non-PII should remain @@ -192,16 +203,34 @@ describe('MCP Server PII Filtering', () => { const result = filterMcpPiiFromSpanData(spanData, false); + // Client info should be filtered expect(result).not.toHaveProperty('client.address'); expect(result).not.toHaveProperty('client.port'); + + // Request arguments should be filtered expect(result).not.toHaveProperty('mcp.request.argument.location'); expect(result).not.toHaveProperty('mcp.request.argument.units'); + + // Specific PII content attributes should be filtered expect(result).not.toHaveProperty('mcp.tool.result.content'); expect(result).not.toHaveProperty('mcp.prompt.result.description'); - expect(result).not.toHaveProperty('mcp.prompt.result.message_content'); + + // Indexed/dynamic result attributes (not in PII_ATTRIBUTES) should remain + expect(result).toHaveProperty('mcp.tool.result.content_count', 1); + expect(result).toHaveProperty('mcp.prompt.result.message_count', 2); + expect(result).toHaveProperty('mcp.prompt.result.0.role', 'user'); + expect(result).toHaveProperty('mcp.prompt.result.0.content', 'Sensitive prompt content'); + expect(result).toHaveProperty('mcp.prompt.result.1.role', 'assistant'); + expect(result).toHaveProperty('mcp.prompt.result.1.content', 'Another sensitive response'); + expect(result).toHaveProperty('mcp.resource.result.content_count', 1); + expect(result).toHaveProperty('mcp.resource.result.uri', 'file:///private/file.txt'); + expect(result).toHaveProperty('mcp.resource.result.content', 'Sensitive resource content'); + + // Other PII attributes should be filtered expect(result).not.toHaveProperty('mcp.logging.message'); expect(result).not.toHaveProperty('mcp.resource.uri'); + // Non-PII attributes should remain expect(result).toHaveProperty('mcp.method.name', 'tools/call'); expect(result).toHaveProperty('mcp.session.id', 'test-session-123'); }); From 1ae03e53e01d3b59be189b3ae16cb565e91414be Mon Sep 17 00:00:00 2001 From: betegon Date: Fri, 1 Aug 2025 17:18:20 +0200 Subject: [PATCH 05/11] ref: keep files under 300 code lines --- .../mcp-server/attributeExtraction.ts | 306 +----------------- .../integrations/mcp-server/correlation.ts | 2 +- .../mcp-server/resultExtraction.ts | 123 +++++++ .../mcp-server/sessionExtraction.ts | 192 +++++++++++ .../src/integrations/mcp-server/transport.ts | 6 +- .../transportInstrumentation.test.ts | 2 +- 6 files changed, 323 insertions(+), 308 deletions(-) create mode 100644 packages/core/src/integrations/mcp-server/resultExtraction.ts create mode 100644 packages/core/src/integrations/mcp-server/sessionExtraction.ts diff --git a/packages/core/src/integrations/mcp-server/attributeExtraction.ts b/packages/core/src/integrations/mcp-server/attributeExtraction.ts index 11b220901033..f5ed09355b74 100644 --- a/packages/core/src/integrations/mcp-server/attributeExtraction.ts +++ b/packages/core/src/integrations/mcp-server/attributeExtraction.ts @@ -1,68 +1,24 @@ /** - * 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_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, - MCP_PROMPT_RESULT_MESSAGE_COUNT_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' }; -} /** * Extracts additional attributes for specific notification types @@ -140,154 +96,7 @@ 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; - 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; - 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 { - const sessionData: Partial = {}; - if (result && typeof result === 'object') { - const resultObj = result as Record; - 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 { - const clientInfo = getClientInfoForTransport(transport); - const attributes: Record = {}; - - 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 { - const serverInfo = getSessionDataForTransport(transport)?.serverInfo; - const attributes: Record = {}; - - 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 { - 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 @@ -315,114 +124,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 { - const attributes: Record = { - [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; - 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; - 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 { - let attributes: Record = {}; - if (typeof result !== 'object' || result === null) return attributes; - - const resultObj = result as Record; - 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; -} - -/** - * Extract prompt result attributes for span instrumentation - * @param result - Prompt execution result - * @returns Attributes extracted from prompt result - */ -export function extractPromptResultAttributes(result: unknown): Record { - const attributes: Record = {}; - if (typeof result !== 'object' || result === null) return attributes; - - const resultObj = result as Record; - - if (typeof resultObj.description === 'string') - attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; - - if (Array.isArray(resultObj.messages)) { - attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length; - - // Extract attributes for each message - const messages = resultObj.messages; - for (const [i, message] of messages.entries()) { - if (typeof message !== 'object' || message === null) continue; - - const messageObj = message as Record; - const prefix = messages.length === 1 ? 'mcp.prompt.result' : `mcp.prompt.result.${i}`; - - const safeSet = (key: string, value: unknown): void => { - if (typeof value === 'string') { - const attrName = messages.length === 1 ? `${prefix}.message_${key}` : `${prefix}.${key}`; - attributes[attrName] = value; - } - }; - - safeSet('role', messageObj.role); - - if (typeof messageObj.content === 'object' && messageObj.content !== null) { - const content = messageObj.content as Record; - if (typeof content.text === 'string') { - const attrName = messages.length === 1 ? `${prefix}.message_content` : `${prefix}.content`; - attributes[attrName] = content.text; - } - } - } - } - - return attributes; -} +// Re-export buildTransportAttributes for spans.ts +export { buildTransportAttributes } from './sessionExtraction'; diff --git a/packages/core/src/integrations/mcp-server/correlation.ts b/packages/core/src/integrations/mcp-server/correlation.ts index 93f9c96280da..39d21f8de033 100644 --- a/packages/core/src/integrations/mcp-server/correlation.ts +++ b/packages/core/src/integrations/mcp-server/correlation.ts @@ -9,8 +9,8 @@ import { getClient } from '../../currentScopes'; import { SPAN_STATUS_ERROR } from '../../tracing'; import type { Span } from '../../types-hoist/span'; -import { extractPromptResultAttributes, extractToolResultAttributes } from './attributeExtraction'; import { filterMcpPiiFromSpanData } from './piiFiltering'; +import { extractPromptResultAttributes, extractToolResultAttributes } from './resultExtraction'; import type { MCPTransport, RequestId, RequestSpanMapValue } from './types'; /** diff --git a/packages/core/src/integrations/mcp-server/resultExtraction.ts b/packages/core/src/integrations/mcp-server/resultExtraction.ts new file mode 100644 index 000000000000..6959b3a0a2f5 --- /dev/null +++ b/packages/core/src/integrations/mcp-server/resultExtraction.ts @@ -0,0 +1,123 @@ +/** + * Result extraction functions for MCP server instrumentation + * + * Handles extraction of attributes from tool and prompt execution results. + */ + +import { + MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE, + MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE, + MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE, +} from './attributes'; + +/** + * 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 { + const attributes: Record = { + [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; + 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; + 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 { + let attributes: Record = {}; + if (typeof result !== 'object' || result === null) return attributes; + + const resultObj = result as Record; + 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; +} + +/** + * Extract prompt result attributes for span instrumentation + * @param result - Prompt execution result + * @returns Attributes extracted from prompt result + */ +export function extractPromptResultAttributes(result: unknown): Record { + const attributes: Record = {}; + if (typeof result !== 'object' || result === null) return attributes; + + const resultObj = result as Record; + + if (typeof resultObj.description === 'string') + attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; + + if (Array.isArray(resultObj.messages)) { + attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length; + + const messages = resultObj.messages; + for (const [i, message] of messages.entries()) { + if (typeof message !== 'object' || message === null) continue; + + const messageObj = message as Record; + const prefix = messages.length === 1 ? 'mcp.prompt.result' : `mcp.prompt.result.${i}`; + + const safeSet = (key: string, value: unknown): void => { + if (typeof value === 'string') { + const attrName = messages.length === 1 ? `${prefix}.message_${key}` : `${prefix}.${key}`; + attributes[attrName] = value; + } + }; + + safeSet('role', messageObj.role); + + if (typeof messageObj.content === 'object' && messageObj.content !== null) { + const content = messageObj.content as Record; + if (typeof content.text === 'string') { + const attrName = messages.length === 1 ? `${prefix}.message_content` : `${prefix}.content`; + attributes[attrName] = content.text; + } + } + } + } + + return attributes; +} diff --git a/packages/core/src/integrations/mcp-server/sessionExtraction.ts b/packages/core/src/integrations/mcp-server/sessionExtraction.ts new file mode 100644 index 000000000000..9265f9d301d4 --- /dev/null +++ b/packages/core/src/integrations/mcp-server/sessionExtraction.ts @@ -0,0 +1,192 @@ +/** + * Session and party info extraction functions for MCP server instrumentation + * + * Handles extraction of client/server info and session data from MCP messages. + */ + +import { + CLIENT_ADDRESS_ATTRIBUTE, + CLIENT_PORT_ATTRIBUTE, + MCP_PROTOCOL_VERSION_ATTRIBUTE, + MCP_SERVER_NAME_ATTRIBUTE, + MCP_SERVER_TITLE_ATTRIBUTE, + MCP_SERVER_VERSION_ATTRIBUTE, + MCP_SESSION_ID_ATTRIBUTE, + MCP_TRANSPORT_ATTRIBUTE, + NETWORK_PROTOCOL_VERSION_ATTRIBUTE, + NETWORK_TRANSPORT_ATTRIBUTE, +} from './attributes'; +import { getClientInfoForTransport, getProtocolVersionForTransport, getSessionDataForTransport } from './sessionManagement'; +import type { ExtraHandlerData, JsonRpcRequest, MCPTransport, PartyInfo, SessionData } from './types'; + +/** + * 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; + 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; + 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 { + const sessionData: Partial = {}; + if (result && typeof result === 'object') { + const resultObj = result as Record; + 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 { + const clientInfo = getClientInfoForTransport(transport); + const attributes: Record = {}; + + 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 { + const serverInfo = getSessionDataForTransport(transport)?.serverInfo; + const attributes: Record = {}; + + 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, + }; +} + +/** + * 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' }; +} + +/** + * 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 { + 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; +} diff --git a/packages/core/src/integrations/mcp-server/transport.ts b/packages/core/src/integrations/mcp-server/transport.ts index 3244ce73e49a..e0950a2774fb 100644 --- a/packages/core/src/integrations/mcp-server/transport.ts +++ b/packages/core/src/integrations/mcp-server/transport.ts @@ -8,12 +8,12 @@ import { getIsolationScope, withIsolationScope } from '../../currentScopes'; import { startInactiveSpan, withActiveSpan } from '../../tracing'; import { fill } from '../../utils/object'; +import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation'; +import { captureError } from './errorCapture'; import { extractSessionDataFromInitializeRequest, extractSessionDataFromInitializeResponse, -} from './attributeExtraction'; -import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation'; -import { captureError } from './errorCapture'; +} from './sessionExtraction'; import { cleanupSessionDataForTransport, storeSessionDataForTransport, diff --git a/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts b/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts index 7f06eb886cdb..008942ac4099 100644 --- a/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts @@ -4,7 +4,7 @@ import { wrapMcpServerWithSentry } from '../../../../src/integrations/mcp-server import { extractSessionDataFromInitializeRequest, extractSessionDataFromInitializeResponse, -} from '../../../../src/integrations/mcp-server/attributeExtraction'; +} from '../../../../src/integrations/mcp-server/sessionExtraction'; import { cleanupSessionDataForTransport, getClientInfoForTransport, From 2c6e5a0c5bbb92d8fc8bec1de6f3470159b4d0d8 Mon Sep 17 00:00:00 2001 From: betegon Date: Fri, 1 Aug 2025 17:34:42 +0200 Subject: [PATCH 06/11] fix lint --- .../src/integrations/mcp-server/attributeExtraction.ts | 10 +--------- .../src/integrations/mcp-server/sessionExtraction.ts | 6 +++++- packages/core/src/integrations/mcp-server/transport.ts | 5 +---- .../lib/integrations/mcp-server/piiFiltering.test.ts | 8 ++++---- 4 files changed, 11 insertions(+), 18 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/attributeExtraction.ts b/packages/core/src/integrations/mcp-server/attributeExtraction.ts index f5ed09355b74..8f1e5a77d94d 100644 --- a/packages/core/src/integrations/mcp-server/attributeExtraction.ts +++ b/packages/core/src/integrations/mcp-server/attributeExtraction.ts @@ -12,13 +12,7 @@ import { MCP_RESOURCE_URI_ATTRIBUTE, } from './attributes'; import { extractTargetInfo, getRequestArguments } from './methodConfig'; -import type { - JsonRpcNotification, - JsonRpcRequest, - McpSpanType, -} from './types'; - - +import type { JsonRpcNotification, JsonRpcRequest, McpSpanType } from './types'; /** * Extracts additional attributes for specific notification types @@ -96,8 +90,6 @@ export function getNotificationAttributes( return attributes; } - - /** * Build type-specific attributes based on message type * @param type - Span type (request or notification) diff --git a/packages/core/src/integrations/mcp-server/sessionExtraction.ts b/packages/core/src/integrations/mcp-server/sessionExtraction.ts index 9265f9d301d4..729d74bb5656 100644 --- a/packages/core/src/integrations/mcp-server/sessionExtraction.ts +++ b/packages/core/src/integrations/mcp-server/sessionExtraction.ts @@ -16,7 +16,11 @@ import { NETWORK_PROTOCOL_VERSION_ATTRIBUTE, NETWORK_TRANSPORT_ATTRIBUTE, } from './attributes'; -import { getClientInfoForTransport, getProtocolVersionForTransport, getSessionDataForTransport } from './sessionManagement'; +import { + getClientInfoForTransport, + getProtocolVersionForTransport, + getSessionDataForTransport, +} from './sessionManagement'; import type { ExtraHandlerData, JsonRpcRequest, MCPTransport, PartyInfo, SessionData } from './types'; /** diff --git a/packages/core/src/integrations/mcp-server/transport.ts b/packages/core/src/integrations/mcp-server/transport.ts index e0950a2774fb..e64b7ab0bdaa 100644 --- a/packages/core/src/integrations/mcp-server/transport.ts +++ b/packages/core/src/integrations/mcp-server/transport.ts @@ -10,10 +10,7 @@ import { startInactiveSpan, withActiveSpan } from '../../tracing'; import { fill } from '../../utils/object'; import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation'; import { captureError } from './errorCapture'; -import { - extractSessionDataFromInitializeRequest, - extractSessionDataFromInitializeResponse, -} from './sessionExtraction'; +import { extractSessionDataFromInitializeRequest, extractSessionDataFromInitializeResponse } from './sessionExtraction'; import { cleanupSessionDataForTransport, storeSessionDataForTransport, diff --git a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts index e79220f7530d..cae9777fc92c 100644 --- a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts @@ -206,15 +206,15 @@ describe('MCP Server PII Filtering', () => { // Client info should be filtered expect(result).not.toHaveProperty('client.address'); expect(result).not.toHaveProperty('client.port'); - + // Request arguments should be filtered expect(result).not.toHaveProperty('mcp.request.argument.location'); expect(result).not.toHaveProperty('mcp.request.argument.units'); - + // Specific PII content attributes should be filtered expect(result).not.toHaveProperty('mcp.tool.result.content'); expect(result).not.toHaveProperty('mcp.prompt.result.description'); - + // Indexed/dynamic result attributes (not in PII_ATTRIBUTES) should remain expect(result).toHaveProperty('mcp.tool.result.content_count', 1); expect(result).toHaveProperty('mcp.prompt.result.message_count', 2); @@ -225,7 +225,7 @@ describe('MCP Server PII Filtering', () => { expect(result).toHaveProperty('mcp.resource.result.content_count', 1); expect(result).toHaveProperty('mcp.resource.result.uri', 'file:///private/file.txt'); expect(result).toHaveProperty('mcp.resource.result.content', 'Sensitive resource content'); - + // Other PII attributes should be filtered expect(result).not.toHaveProperty('mcp.logging.message'); expect(result).not.toHaveProperty('mcp.resource.uri'); From d1e9fdc14e1361ebb7c6fd3e7378a85d5ff8a6eb Mon Sep 17 00:00:00 2001 From: betegon Date: Mon, 4 Aug 2025 17:02:49 +0200 Subject: [PATCH 07/11] avoid inline conditionals --- .../integrations/mcp-server/correlation.ts | 4 ++- .../mcp-server/resultExtraction.ts | 29 ++++++++++++++----- .../mcp-server/sessionExtraction.ts | 16 +++++++--- .../mcp-server/sessionManagement.ts | 4 ++- 4 files changed, 40 insertions(+), 13 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/correlation.ts b/packages/core/src/integrations/mcp-server/correlation.ts index 39d21f8de033..7a73f63f64e3 100644 --- a/packages/core/src/integrations/mcp-server/correlation.ts +++ b/packages/core/src/integrations/mcp-server/correlation.ts @@ -90,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; diff --git a/packages/core/src/integrations/mcp-server/resultExtraction.ts b/packages/core/src/integrations/mcp-server/resultExtraction.ts index 6959b3a0a2f5..0c055cc7345d 100644 --- a/packages/core/src/integrations/mcp-server/resultExtraction.ts +++ b/packages/core/src/integrations/mcp-server/resultExtraction.ts @@ -22,13 +22,17 @@ function buildAllContentItemAttributes(content: unknown[]): Record; 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; + if (typeof value === 'string') { + attributes[`${prefix}.${key}`] = value; + } }; safeSet('content_type', contentItem.type); @@ -39,7 +43,11 @@ function buildAllContentItemAttributes(content: unknown[]): Record maxLength ? `${text.slice(0, maxLength - 3)}...` : text; + if (text.length > maxLength) { + attributes[`${prefix}.content`] = `${text.slice(0, maxLength - 3)}...`; + } else { + attributes[`${prefix}.content`] = text; + } } if (typeof contentItem.data === 'string') { @@ -64,7 +72,9 @@ function buildAllContentItemAttributes(content: unknown[]): Record { let attributes: Record = {}; - if (typeof result !== 'object' || result === null) return attributes; + if (typeof result !== 'object' || result === null) { + return attributes; + } const resultObj = result as Record; if (typeof resultObj.isError === 'boolean') { @@ -83,19 +93,24 @@ export function extractToolResultAttributes(result: unknown): Record { const attributes: Record = {}; - if (typeof result !== 'object' || result === null) return attributes; + if (typeof result !== 'object' || result === null) { + return attributes; + } const resultObj = result as Record; - if (typeof resultObj.description === 'string') + if (typeof resultObj.description === 'string') { attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; + } if (Array.isArray(resultObj.messages)) { attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length; const messages = resultObj.messages; for (const [i, message] of messages.entries()) { - if (typeof message !== 'object' || message === null) continue; + if (typeof message !== 'object' || message === null) { + continue; + } const messageObj = message as Record; const prefix = messages.length === 1 ? 'mcp.prompt.result' : `mcp.prompt.result.${i}`; diff --git a/packages/core/src/integrations/mcp-server/sessionExtraction.ts b/packages/core/src/integrations/mcp-server/sessionExtraction.ts index 729d74bb5656..7e068cc10732 100644 --- a/packages/core/src/integrations/mcp-server/sessionExtraction.ts +++ b/packages/core/src/integrations/mcp-server/sessionExtraction.ts @@ -33,9 +33,15 @@ function extractPartyInfo(obj: unknown): PartyInfo { if (obj && typeof obj === 'object' && obj !== null) { const source = obj as Record; - 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; + 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; @@ -69,7 +75,9 @@ export function extractSessionDataFromInitializeResponse(result: unknown): Parti const sessionData: Partial = {}; if (result && typeof result === 'object') { const resultObj = result as Record; - if (typeof resultObj.protocolVersion === 'string') sessionData.protocolVersion = resultObj.protocolVersion; + if (typeof resultObj.protocolVersion === 'string') { + sessionData.protocolVersion = resultObj.protocolVersion; + } if (resultObj.serverInfo) { sessionData.serverInfo = extractPartyInfo(resultObj.serverInfo); } diff --git a/packages/core/src/integrations/mcp-server/sessionManagement.ts b/packages/core/src/integrations/mcp-server/sessionManagement.ts index 99ba2e0d8806..9d9c8b48f27d 100644 --- a/packages/core/src/integrations/mcp-server/sessionManagement.ts +++ b/packages/core/src/integrations/mcp-server/sessionManagement.ts @@ -16,7 +16,9 @@ const transportToSessionData = new WeakMap(); * @param sessionData - Session data to store */ export function storeSessionDataForTransport(transport: MCPTransport, sessionData: SessionData): void { - if (transport.sessionId) transportToSessionData.set(transport, sessionData); + if (transport.sessionId) { + transportToSessionData.set(transport, sessionData); + } } /** From a0fb99e878bd4384a3c5cb6a0f110ab857029bcd Mon Sep 17 00:00:00 2001 From: betegon Date: Mon, 4 Aug 2025 18:41:47 +0200 Subject: [PATCH 08/11] use type narrowing instead of casting --- .../mcp-server/resultExtraction.ts | 63 +++++++++---------- .../mcp-server/sessionExtraction.ts | 38 ++++++----- .../src/integrations/mcp-server/transport.ts | 7 +-- .../src/integrations/mcp-server/validation.ts | 9 +++ 4 files changed, 59 insertions(+), 58 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/resultExtraction.ts b/packages/core/src/integrations/mcp-server/resultExtraction.ts index 0c055cc7345d..79c790ab40de 100644 --- a/packages/core/src/integrations/mcp-server/resultExtraction.ts +++ b/packages/core/src/integrations/mcp-server/resultExtraction.ts @@ -10,6 +10,7 @@ import { MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE, MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE, } from './attributes'; +import { isValidContentItem } from './validation'; /** * Build attributes for tool result content items @@ -22,11 +23,10 @@ function buildAllContentItemAttributes(content: unknown[]): Record; const prefix = content.length === 1 ? 'mcp.tool.result' : `mcp.tool.result.${i}`; const safeSet = (key: string, value: unknown): void => { @@ -35,13 +35,13 @@ function buildAllContentItemAttributes(content: unknown[]): Record maxLength) { attributes[`${prefix}.content`] = `${text.slice(0, maxLength - 3)}...`; @@ -50,15 +50,14 @@ function buildAllContentItemAttributes(content: unknown[]): Record; - safeSet('resource_uri', res.uri); - safeSet('resource_mime_type', res.mimeType); + const resource = item.resource; + if (isValidContentItem(resource)) { + safeSet('resource_uri', resource.uri); + safeSet('resource_mime_type', resource.mimeType); } } @@ -72,16 +71,15 @@ function buildAllContentItemAttributes(content: unknown[]): Record { let attributes: Record = {}; - if (typeof result !== 'object' || result === null) { + if (!isValidContentItem(result)) { return attributes; } - const resultObj = result as Record; - if (typeof resultObj.isError === 'boolean') { - attributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] = resultObj.isError; + if (typeof result.isError === 'boolean') { + attributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] = result.isError; } - if (Array.isArray(resultObj.content)) { - attributes = { ...attributes, ...buildAllContentItemAttributes(resultObj.content) }; + if (Array.isArray(result.content)) { + attributes = { ...attributes, ...buildAllContentItemAttributes(result.content) }; } return attributes; } @@ -93,26 +91,23 @@ export function extractToolResultAttributes(result: unknown): Record { const attributes: Record = {}; - if (typeof result !== 'object' || result === null) { + if (!isValidContentItem(result)) { return attributes; } - const resultObj = result as Record; - - if (typeof resultObj.description === 'string') { - attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description; + if (typeof result.description === 'string') { + attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = result.description; } - if (Array.isArray(resultObj.messages)) { - attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length; + if (Array.isArray(result.messages)) { + attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = result.messages.length; - const messages = resultObj.messages; + const messages = result.messages; for (const [i, message] of messages.entries()) { - if (typeof message !== 'object' || message === null) { + if (!isValidContentItem(message)) { continue; } - const messageObj = message as Record; const prefix = messages.length === 1 ? 'mcp.prompt.result' : `mcp.prompt.result.${i}`; const safeSet = (key: string, value: unknown): void => { @@ -122,10 +117,10 @@ export function extractPromptResultAttributes(result: unknown): Record; + if (isValidContentItem(message.content)) { + const content = message.content; if (typeof content.text === 'string') { const attrName = messages.length === 1 ? `${prefix}.message_content` : `${prefix}.content`; attributes[attrName] = content.text; diff --git a/packages/core/src/integrations/mcp-server/sessionExtraction.ts b/packages/core/src/integrations/mcp-server/sessionExtraction.ts index 7e068cc10732..90e235d4e544 100644 --- a/packages/core/src/integrations/mcp-server/sessionExtraction.ts +++ b/packages/core/src/integrations/mcp-server/sessionExtraction.ts @@ -22,6 +22,7 @@ import { getSessionDataForTransport, } from './sessionManagement'; import type { ExtraHandlerData, JsonRpcRequest, MCPTransport, PartyInfo, SessionData } from './types'; +import { isValidContentItem } from './validation'; /** * Extracts and validates PartyInfo from an unknown object @@ -31,16 +32,15 @@ import type { ExtraHandlerData, JsonRpcRequest, MCPTransport, PartyInfo, Session function extractPartyInfo(obj: unknown): PartyInfo { const partyInfo: PartyInfo = {}; - if (obj && typeof obj === 'object' && obj !== null) { - const source = obj as Record; - if (typeof source.name === 'string') { - partyInfo.name = source.name; + if (isValidContentItem(obj)) { + if (typeof obj.name === 'string') { + partyInfo.name = obj.name; } - if (typeof source.title === 'string') { - partyInfo.title = source.title; + if (typeof obj.title === 'string') { + partyInfo.title = obj.title; } - if (typeof source.version === 'string') { - partyInfo.version = source.version; + if (typeof obj.version === 'string') { + partyInfo.version = obj.version; } } @@ -54,13 +54,12 @@ function extractPartyInfo(obj: unknown): PartyInfo { */ 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; - if (typeof params.protocolVersion === 'string') { - sessionData.protocolVersion = params.protocolVersion; + if (isValidContentItem(request.params)) { + if (typeof request.params.protocolVersion === 'string') { + sessionData.protocolVersion = request.params.protocolVersion; } - if (params.clientInfo) { - sessionData.clientInfo = extractPartyInfo(params.clientInfo); + if (request.params.clientInfo) { + sessionData.clientInfo = extractPartyInfo(request.params.clientInfo); } } return sessionData; @@ -73,13 +72,12 @@ export function extractSessionDataFromInitializeRequest(request: JsonRpcRequest) */ export function extractSessionDataFromInitializeResponse(result: unknown): Partial { const sessionData: Partial = {}; - if (result && typeof result === 'object') { - const resultObj = result as Record; - if (typeof resultObj.protocolVersion === 'string') { - sessionData.protocolVersion = resultObj.protocolVersion; + if (isValidContentItem(result)) { + if (typeof result.protocolVersion === 'string') { + sessionData.protocolVersion = result.protocolVersion; } - if (resultObj.serverInfo) { - sessionData.serverInfo = extractPartyInfo(resultObj.serverInfo); + if (result.serverInfo) { + sessionData.serverInfo = extractPartyInfo(result.serverInfo); } } return sessionData; diff --git a/packages/core/src/integrations/mcp-server/transport.ts b/packages/core/src/integrations/mcp-server/transport.ts index e64b7ab0bdaa..6943ac3e8850 100644 --- a/packages/core/src/integrations/mcp-server/transport.ts +++ b/packages/core/src/integrations/mcp-server/transport.ts @@ -18,7 +18,7 @@ import { } from './sessionManagement'; import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans'; import type { ExtraHandlerData, MCPTransport } from './types'; -import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse } from './validation'; +import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse, isValidContentItem } from './validation'; /** * Wraps transport.onmessage to create spans for incoming messages. @@ -90,9 +90,8 @@ export function wrapTransportSend(transport: MCPTransport): void { captureJsonRpcErrorResponse(message.error); } - if (message.result && typeof message.result === 'object') { - const result = message.result as Record; - if (result.protocolVersion || result.serverInfo) { + if (isValidContentItem(message.result)) { + if (message.result.protocolVersion || message.result.serverInfo) { try { const serverData = extractSessionDataFromInitializeResponse(message.result); updateSessionDataForTransport(this, serverData); diff --git a/packages/core/src/integrations/mcp-server/validation.ts b/packages/core/src/integrations/mcp-server/validation.ts index 21d257c01aeb..9ed21b290728 100644 --- a/packages/core/src/integrations/mcp-server/validation.ts +++ b/packages/core/src/integrations/mcp-server/validation.ts @@ -75,3 +75,12 @@ export function validateMcpServerInstance(instance: unknown): boolean { DEBUG_BUILD && debug.warn('Did not patch MCP server. Interface is incompatible.'); return false; } + +/** + * Check if the item is a valid content item + * @param item - The item to check + * @returns True if the item is a valid content item, false otherwise + */ +export function isValidContentItem(item: unknown): item is Record { + return item != null && typeof item === 'object'; +} From 105e8b50aabdafd3a12c19135113a3583f7e0989 Mon Sep 17 00:00:00 2001 From: betegon Date: Mon, 4 Aug 2025 19:02:07 +0200 Subject: [PATCH 09/11] remove length limit --- .../core/src/integrations/mcp-server/resultExtraction.ts | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/resultExtraction.ts b/packages/core/src/integrations/mcp-server/resultExtraction.ts index 79c790ab40de..766b29a718d0 100644 --- a/packages/core/src/integrations/mcp-server/resultExtraction.ts +++ b/packages/core/src/integrations/mcp-server/resultExtraction.ts @@ -41,13 +41,7 @@ function buildAllContentItemAttributes(content: unknown[]): Record maxLength) { - attributes[`${prefix}.content`] = `${text.slice(0, maxLength - 3)}...`; - } else { - attributes[`${prefix}.content`] = text; - } + attributes[`${prefix}.content`] = item.text; } if (typeof item.data === 'string') { From 38a6ea50fd3d8e0abe4ee54a76fd4d6178763ba8 Mon Sep 17 00:00:00 2001 From: betegon Date: Mon, 4 Aug 2025 19:11:25 +0200 Subject: [PATCH 10/11] avoid spreads --- .../src/integrations/mcp-server/resultExtraction.ts | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/resultExtraction.ts b/packages/core/src/integrations/mcp-server/resultExtraction.ts index 766b29a718d0..34dc2be9d09c 100644 --- a/packages/core/src/integrations/mcp-server/resultExtraction.ts +++ b/packages/core/src/integrations/mcp-server/resultExtraction.ts @@ -17,7 +17,7 @@ import { isValidContentItem } from './validation'; * @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 { +function buildAllContentItemAttributes(content: unknown[]): Record { const attributes: Record = { [MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE]: content.length, }; @@ -64,17 +64,16 @@ function buildAllContentItemAttributes(content: unknown[]): Record { - let attributes: Record = {}; if (!isValidContentItem(result)) { - return attributes; + return {}; } + const attributes = Array.isArray(result.content) ? buildAllContentItemAttributes(result.content) : {}; + if (typeof result.isError === 'boolean') { attributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] = result.isError; } - if (Array.isArray(result.content)) { - attributes = { ...attributes, ...buildAllContentItemAttributes(result.content) }; - } + return attributes; } From 41556c8f49b9369461b38bac474128bca710bfad Mon Sep 17 00:00:00 2001 From: betegon Date: Wed, 6 Aug 2025 11:53:43 +0200 Subject: [PATCH 11/11] take into account result list attributes as PII --- .../src/integrations/mcp-server/attributes.ts | 6 ++++ .../integrations/mcp-server/piiFiltering.ts | 30 +++++++++++++++++-- .../mcp-server/piiFiltering.test.ts | 13 ++++---- 3 files changed, 42 insertions(+), 7 deletions(-) diff --git a/packages/core/src/integrations/mcp-server/attributes.ts b/packages/core/src/integrations/mcp-server/attributes.ts index faad50052347..273bdbdb9560 100644 --- a/packages/core/src/integrations/mcp-server/attributes.ts +++ b/packages/core/src/integrations/mcp-server/attributes.ts @@ -76,6 +76,9 @@ 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 // ============================================================================= @@ -92,6 +95,9 @@ export const MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE = 'mcp.prompt.result.messa /** 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 // ============================================================================= diff --git a/packages/core/src/integrations/mcp-server/piiFiltering.ts b/packages/core/src/integrations/mcp-server/piiFiltering.ts index 0161b926cb23..ff801cbf2a1e 100644 --- a/packages/core/src/integrations/mcp-server/piiFiltering.ts +++ b/packages/core/src/integrations/mcp-server/piiFiltering.ts @@ -11,9 +11,11 @@ import { MCP_LOGGING_MESSAGE_ATTRIBUTE, MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE, + MCP_PROMPT_RESULT_PREFIX, MCP_REQUEST_ARGUMENT, MCP_RESOURCE_URI_ATTRIBUTE, MCP_TOOL_RESULT_CONTENT_ATTRIBUTE, + MCP_TOOL_RESULT_PREFIX, } from './attributes'; /** @@ -31,11 +33,35 @@ const PII_ATTRIBUTES = new Set([ ]); /** - * Checks if an attribute key should be considered PII + * Checks if an attribute key should be considered PII. + * + * Returns true for: + * - Explicit PII attributes (client.address, client.port, mcp.logging.message, etc.) + * - All request arguments (mcp.request.argument.*) + * - Tool and prompt result content (mcp.tool.result.*, mcp.prompt.result.*) except metadata + * + * Preserves metadata attributes ending with _count, _error, or .is_error as they don't contain sensitive data. + * + * @param key - Attribute key to evaluate + * @returns true if the attribute should be filtered out (is PII), false if it should be preserved * @internal */ function isPiiAttribute(key: string): boolean { - return PII_ATTRIBUTES.has(key) || key.startsWith(`${MCP_REQUEST_ARGUMENT}.`); + if (PII_ATTRIBUTES.has(key)) { + return true; + } + + if (key.startsWith(`${MCP_REQUEST_ARGUMENT}.`)) { + return true; + } + + if (key.startsWith(`${MCP_TOOL_RESULT_PREFIX}.`) || key.startsWith(`${MCP_PROMPT_RESULT_PREFIX}.`)) { + if (!key.endsWith('_count') && !key.endsWith('_error') && !key.endsWith('.is_error')) { + return true; + } + } + + return false; } /** diff --git a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts index cae9777fc92c..a86ccbd534d0 100644 --- a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts @@ -215,13 +215,16 @@ describe('MCP Server PII Filtering', () => { expect(result).not.toHaveProperty('mcp.tool.result.content'); expect(result).not.toHaveProperty('mcp.prompt.result.description'); - // Indexed/dynamic result attributes (not in PII_ATTRIBUTES) should remain + // Count attributes should remain as they don't contain sensitive content expect(result).toHaveProperty('mcp.tool.result.content_count', 1); expect(result).toHaveProperty('mcp.prompt.result.message_count', 2); - expect(result).toHaveProperty('mcp.prompt.result.0.role', 'user'); - expect(result).toHaveProperty('mcp.prompt.result.0.content', 'Sensitive prompt content'); - expect(result).toHaveProperty('mcp.prompt.result.1.role', 'assistant'); - expect(result).toHaveProperty('mcp.prompt.result.1.content', 'Another sensitive response'); + + // All tool and prompt result content should be filtered (including indexed attributes) + expect(result).not.toHaveProperty('mcp.prompt.result.0.role'); + expect(result).not.toHaveProperty('mcp.prompt.result.0.content'); + expect(result).not.toHaveProperty('mcp.prompt.result.1.role'); + expect(result).not.toHaveProperty('mcp.prompt.result.1.content'); + expect(result).toHaveProperty('mcp.resource.result.content_count', 1); expect(result).toHaveProperty('mcp.resource.result.uri', 'file:///private/file.txt'); expect(result).toHaveProperty('mcp.resource.result.content', 'Sensitive resource content');