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.

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
39 changes: 39 additions & 0 deletions packages/core/src/integrations/mcp-server/attributeExtraction.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -377,3 +381,38 @@ export function extractToolResultAttributes(result: unknown): Record<string, str
}
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<string, string | number | boolean> {
const attributes: Record<string, string | number | boolean> = {};
if (typeof result !== 'object' || result === null) return attributes;

const resultObj = result as Record<string, unknown>;

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<string, unknown>;

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<string, unknown>;
if (typeof content.text === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text;
}
}
}
}

return attributes;
}
16 changes: 16 additions & 0 deletions packages/core/src/integrations/mcp-server/attributes.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
// =============================================================================
Expand Down
9 changes: 8 additions & 1 deletion packages/core/src/integrations/mcp-server/correlation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';

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 Down
4 changes: 4 additions & 0 deletions packages/core/src/integrations/mcp-server/piiFiltering.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
]);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ describe('MCP Server PII Filtering', () => {
setAttributes: vi.fn(),
setStatus: vi.fn(),
end: vi.fn(),
} as any;
} as unknown as ReturnType<typeof tracingModule.startInactiveSpan>;
startInactiveSpanSpy.mockReturnValueOnce(mockSpan);

const toolCallRequest = {
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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');

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -434,5 +434,77 @@ 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<typeof tracingModule.startInactiveSpan>,
);

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();
});
});
});