Skip to content

Commit f1fbea0

Browse files
betegonbetegon
committed
Capture prompt results in prompt function calls
1 parent eda7e33 commit f1fbea0

File tree

6 files changed

+153
-2
lines changed

6 files changed

+153
-2
lines changed

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@ import {
1010
MCP_LOGGING_LEVEL_ATTRIBUTE,
1111
MCP_LOGGING_LOGGER_ATTRIBUTE,
1212
MCP_LOGGING_MESSAGE_ATTRIBUTE,
13+
MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE,
14+
MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE,
15+
MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE,
16+
MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE,
1317
MCP_PROTOCOL_VERSION_ATTRIBUTE,
1418
MCP_REQUEST_ID_ATTRIBUTE,
1519
MCP_RESOURCE_URI_ATTRIBUTE,
@@ -377,3 +381,43 @@ export function extractToolResultAttributes(result: unknown): Record<string, str
377381
}
378382
return attributes;
379383
}
384+
385+
/**
386+
* Extract prompt result attributes for span instrumentation
387+
* @param result - Prompt execution result
388+
* @returns Attributes extracted from prompt result
389+
*/
390+
export function extractPromptResultAttributes(result: unknown): Record<string, string | number | boolean> {
391+
const attributes: Record<string, string | number | boolean> = {};
392+
if (typeof result !== 'object' || result === null) return attributes;
393+
394+
const resultObj = result as Record<string, unknown>;
395+
396+
if (typeof resultObj.description === 'string') {
397+
attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = resultObj.description;
398+
}
399+
400+
if (Array.isArray(resultObj.messages)) {
401+
attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length;
402+
403+
if (resultObj.messages.length > 0) {
404+
const message = resultObj.messages[0];
405+
if (typeof message === 'object' && message !== null) {
406+
const messageObj = message as Record<string, unknown>;
407+
408+
if (typeof messageObj.role === 'string') {
409+
attributes[MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE] = messageObj.role;
410+
}
411+
412+
if (typeof messageObj.content === 'object' && messageObj.content !== null) {
413+
const content = messageObj.content as Record<string, unknown>;
414+
if (typeof content.text === 'string') {
415+
attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text;
416+
}
417+
}
418+
}
419+
}
420+
}
421+
422+
return attributes;
423+
}

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

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

79+
// =============================================================================
80+
// PROMPT RESULT ATTRIBUTES
81+
// =============================================================================
82+
83+
/** Description of the prompt result */
84+
export const MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE = 'mcp.prompt.result.description';
85+
86+
/** Number of messages in the prompt result */
87+
export const MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE = 'mcp.prompt.result.message_count';
88+
89+
/** Role of the prompt message (for single message prompts) */
90+
export const MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE = 'mcp.prompt.result.message_role';
91+
92+
/** Content of the prompt message (for single message prompts) */
93+
export const MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE = 'mcp.prompt.result.message_content';
94+
7995
// =============================================================================
8096
// REQUEST ARGUMENT ATTRIBUTES
8197
// =============================================================================

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

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

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

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

7481
span.end();

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import {
99
CLIENT_ADDRESS_ATTRIBUTE,
1010
CLIENT_PORT_ATTRIBUTE,
1111
MCP_LOGGING_MESSAGE_ATTRIBUTE,
12+
MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE,
13+
MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE,
1214
MCP_REQUEST_ARGUMENT,
1315
MCP_RESOURCE_URI_ATTRIBUTE,
1416
MCP_TOOL_RESULT_CONTENT_ATTRIBUTE,
@@ -22,6 +24,8 @@ const PII_ATTRIBUTES = new Set([
2224
CLIENT_ADDRESS_ATTRIBUTE,
2325
CLIENT_PORT_ATTRIBUTE,
2426
MCP_LOGGING_MESSAGE_ATTRIBUTE,
27+
MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE,
28+
MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE,
2529
MCP_RESOURCE_URI_ATTRIBUTE,
2630
MCP_TOOL_RESULT_CONTENT_ATTRIBUTE,
2731
]);

‎packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ describe('MCP Server PII Filtering', () => {
124124
setAttributes: vi.fn(),
125125
setStatus: vi.fn(),
126126
end: vi.fn(),
127-
} as any;
127+
} as unknown as ReturnType<typeof tracingModule.startInactiveSpan>;
128128
startInactiveSpanSpy.mockReturnValueOnce(mockSpan);
129129

130130
const toolCallRequest = {
@@ -163,6 +163,8 @@ describe('MCP Server PII Filtering', () => {
163163
'client.port': 54321,
164164
'mcp.request.argument.location': '"San Francisco"',
165165
'mcp.tool.result.content': 'Weather data: 18°C',
166+
'mcp.prompt.result.description': 'Code review prompt for sensitive analysis',
167+
'mcp.prompt.result.message_content': 'Please review this confidential code.',
166168
'mcp.logging.message': 'User requested weather',
167169
'mcp.resource.uri': 'file:///private/docs/secret.txt',
168170
'mcp.method.name': 'tools/call', // Non-PII should remain
@@ -180,6 +182,8 @@ describe('MCP Server PII Filtering', () => {
180182
'mcp.request.argument.location': '"San Francisco"',
181183
'mcp.request.argument.units': '"celsius"',
182184
'mcp.tool.result.content': 'Weather data: 18°C',
185+
'mcp.prompt.result.description': 'Code review prompt for sensitive analysis',
186+
'mcp.prompt.result.message_content': 'Please review this confidential code.',
183187
'mcp.logging.message': 'User requested weather',
184188
'mcp.resource.uri': 'file:///private/docs/secret.txt',
185189
'mcp.method.name': 'tools/call', // Non-PII should remain
@@ -193,6 +197,8 @@ describe('MCP Server PII Filtering', () => {
193197
expect(result).not.toHaveProperty('mcp.request.argument.location');
194198
expect(result).not.toHaveProperty('mcp.request.argument.units');
195199
expect(result).not.toHaveProperty('mcp.tool.result.content');
200+
expect(result).not.toHaveProperty('mcp.prompt.result.description');
201+
expect(result).not.toHaveProperty('mcp.prompt.result.message_content');
196202
expect(result).not.toHaveProperty('mcp.logging.message');
197203
expect(result).not.toHaveProperty('mcp.resource.uri');
198204

‎packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -434,5 +434,79 @@ describe('MCP Server Semantic Conventions', () => {
434434
expect(setStatusSpy).not.toHaveBeenCalled();
435435
expect(endSpy).toHaveBeenCalled();
436436
});
437+
438+
it('should instrument prompt call results and complete span with enriched attributes', async () => {
439+
await wrappedMcpServer.connect(mockTransport);
440+
441+
const setAttributesSpy = vi.fn();
442+
const setStatusSpy = vi.fn();
443+
const endSpy = vi.fn();
444+
const mockSpan = {
445+
setAttributes: setAttributesSpy,
446+
setStatus: setStatusSpy,
447+
end: endSpy,
448+
};
449+
startInactiveSpanSpy.mockReturnValueOnce(
450+
mockSpan as unknown as ReturnType<typeof tracingModule.startInactiveSpan>,
451+
);
452+
453+
const promptCallRequest = {
454+
jsonrpc: '2.0',
455+
method: 'prompts/get',
456+
id: 'req-prompt-result',
457+
params: {
458+
name: 'code-review',
459+
arguments: { language: 'typescript', complexity: 'high' },
460+
},
461+
};
462+
463+
mockTransport.onmessage?.(promptCallRequest, {});
464+
465+
expect(startInactiveSpanSpy).toHaveBeenCalledWith(
466+
expect.objectContaining({
467+
name: 'prompts/get code-review',
468+
op: 'mcp.server',
469+
forceTransaction: true,
470+
attributes: expect.objectContaining({
471+
'mcp.method.name': 'prompts/get',
472+
'mcp.prompt.name': 'code-review',
473+
'mcp.request.id': 'req-prompt-result',
474+
}),
475+
}),
476+
);
477+
478+
const promptResponse = {
479+
jsonrpc: '2.0',
480+
id: 'req-prompt-result',
481+
result: {
482+
description: 'Code review prompt for TypeScript with high complexity analysis',
483+
messages: [
484+
{
485+
role: 'user',
486+
content: {
487+
type: 'text',
488+
text: 'Please review this TypeScript code for complexity and best practices.',
489+
},
490+
},
491+
],
492+
},
493+
};
494+
495+
mockTransport.send?.(promptResponse);
496+
497+
expect(setAttributesSpy).toHaveBeenCalledWith(
498+
expect.objectContaining({
499+
'mcp.prompt.result.description': 'Code review prompt for TypeScript with high complexity analysis',
500+
'mcp.prompt.result.message_count': 1,
501+
'mcp.prompt.result.message_role': 'user',
502+
'mcp.prompt.result.message_content': 'Please review this TypeScript code for complexity and best practices.',
503+
}),
504+
);
505+
506+
expect(setStatusSpy).not.toHaveBeenCalled();
507+
expect(endSpy).toHaveBeenCalled();
508+
});
509+
510+
437511
});
438512
});

0 commit comments

Comments
 (0)