Skip to content

Commit 972d76d

Browse files
betegonbetegon
committed
capture ALL messages in result.
1 parent 18d74f1 commit 972d76d

File tree

3 files changed

+54
-15
lines changed

3 files changed

+54
-15
lines changed

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

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,7 @@ import {
1111
MCP_LOGGING_LOGGER_ATTRIBUTE,
1212
MCP_LOGGING_MESSAGE_ATTRIBUTE,
1313
MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE,
14-
MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE,
1514
MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE,
16-
MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE,
1715
MCP_PROTOCOL_VERSION_ATTRIBUTE,
1816
MCP_REQUEST_ID_ATTRIBUTE,
1917
MCP_RESOURCE_URI_ATTRIBUTE,
@@ -399,16 +397,28 @@ export function extractPromptResultAttributes(result: unknown): Record<string, s
399397
if (Array.isArray(resultObj.messages)) {
400398
attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = resultObj.messages.length;
401399

402-
if (resultObj.messages.length > 0) {
403-
const message = resultObj.messages[0];
404-
if (typeof message === 'object' && message !== null) {
405-
const messageObj = message as Record<string, unknown>;
400+
// Extract attributes for each message
401+
const messages = resultObj.messages;
402+
for (const [i, message] of messages.entries()) {
403+
if (typeof message !== 'object' || message === null) continue;
406404

407-
if (typeof messageObj.role === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE] = messageObj.role;
405+
const messageObj = message as Record<string, unknown>;
406+
const prefix = messages.length === 1 ? 'mcp.prompt.result' : `mcp.prompt.result.${i}`;
408407

409-
if (typeof messageObj.content === 'object' && messageObj.content !== null) {
410-
const content = messageObj.content as Record<string, unknown>;
411-
if (typeof content.text === 'string') attributes[MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE] = content.text;
408+
const safeSet = (key: string, value: unknown): void => {
409+
if (typeof value === 'string') {
410+
const attrName = messages.length === 1 ? `${prefix}.message_${key}` : `${prefix}.${key}`;
411+
attributes[attrName] = value;
412+
}
413+
};
414+
415+
safeSet('role', messageObj.role);
416+
417+
if (typeof messageObj.content === 'object' && messageObj.content !== null) {
418+
const content = messageObj.content as Record<string, unknown>;
419+
if (typeof content.text === 'string') {
420+
const attrName = messages.length === 1 ? `${prefix}.message_content` : `${prefix}.content`;
421+
attributes[attrName] = content.text;
412422
}
413423
}
414424
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,10 +86,10 @@ export const MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE = 'mcp.prompt.result.descri
8686
/** Number of messages in the prompt result */
8787
export const MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE = 'mcp.prompt.result.message_count';
8888

89-
/** Role of the prompt message (for single message prompts) */
89+
/** Role of the message in the prompt result (for single message results) */
9090
export const MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE = 'mcp.prompt.result.message_role';
9191

92-
/** Content of the prompt message (for single message prompts) */
92+
/** Content of the message in the prompt result (for single message results) */
9393
export const MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE = 'mcp.prompt.result.message_content';
9494

9595
// =============================================================================

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

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ describe('MCP Server PII Filtering', () => {
147147

148148
mockTransport.send?.(toolResponse);
149149

150-
// Tool result content should be filtered out
150+
// Tool result content should be filtered out, but metadata should remain
151151
const setAttributesCall = mockSpan.setAttributes.mock.calls[0]?.[0];
152152
expect(setAttributesCall).toBeDefined();
153153
expect(setAttributesCall).not.toHaveProperty('mcp.tool.result.content');
@@ -163,8 +163,11 @@ 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.tool.result.content_count': 1,
166167
'mcp.prompt.result.description': 'Code review prompt for sensitive analysis',
167168
'mcp.prompt.result.message_content': 'Please review this confidential code.',
169+
'mcp.prompt.result.message_count': 1,
170+
'mcp.resource.result.content': 'Sensitive resource content',
168171
'mcp.logging.message': 'User requested weather',
169172
'mcp.resource.uri': 'file:///private/docs/secret.txt',
170173
'mcp.method.name': 'tools/call', // Non-PII should remain
@@ -182,8 +185,16 @@ describe('MCP Server PII Filtering', () => {
182185
'mcp.request.argument.location': '"San Francisco"',
183186
'mcp.request.argument.units': '"celsius"',
184187
'mcp.tool.result.content': 'Weather data: 18°C',
188+
'mcp.tool.result.content_count': 1,
185189
'mcp.prompt.result.description': 'Code review prompt for sensitive analysis',
186-
'mcp.prompt.result.message_content': 'Please review this confidential code.',
190+
'mcp.prompt.result.message_count': 2,
191+
'mcp.prompt.result.0.role': 'user',
192+
'mcp.prompt.result.0.content': 'Sensitive prompt content',
193+
'mcp.prompt.result.1.role': 'assistant',
194+
'mcp.prompt.result.1.content': 'Another sensitive response',
195+
'mcp.resource.result.content_count': 1,
196+
'mcp.resource.result.uri': 'file:///private/file.txt',
197+
'mcp.resource.result.content': 'Sensitive resource content',
187198
'mcp.logging.message': 'User requested weather',
188199
'mcp.resource.uri': 'file:///private/docs/secret.txt',
189200
'mcp.method.name': 'tools/call', // Non-PII should remain
@@ -192,16 +203,34 @@ describe('MCP Server PII Filtering', () => {
192203

193204
const result = filterMcpPiiFromSpanData(spanData, false);
194205

206+
// Client info should be filtered
195207
expect(result).not.toHaveProperty('client.address');
196208
expect(result).not.toHaveProperty('client.port');
209+
210+
// Request arguments should be filtered
197211
expect(result).not.toHaveProperty('mcp.request.argument.location');
198212
expect(result).not.toHaveProperty('mcp.request.argument.units');
213+
214+
// Specific PII content attributes should be filtered
199215
expect(result).not.toHaveProperty('mcp.tool.result.content');
200216
expect(result).not.toHaveProperty('mcp.prompt.result.description');
201-
expect(result).not.toHaveProperty('mcp.prompt.result.message_content');
217+
218+
// Indexed/dynamic result attributes (not in PII_ATTRIBUTES) should remain
219+
expect(result).toHaveProperty('mcp.tool.result.content_count', 1);
220+
expect(result).toHaveProperty('mcp.prompt.result.message_count', 2);
221+
expect(result).toHaveProperty('mcp.prompt.result.0.role', 'user');
222+
expect(result).toHaveProperty('mcp.prompt.result.0.content', 'Sensitive prompt content');
223+
expect(result).toHaveProperty('mcp.prompt.result.1.role', 'assistant');
224+
expect(result).toHaveProperty('mcp.prompt.result.1.content', 'Another sensitive response');
225+
expect(result).toHaveProperty('mcp.resource.result.content_count', 1);
226+
expect(result).toHaveProperty('mcp.resource.result.uri', 'file:///private/file.txt');
227+
expect(result).toHaveProperty('mcp.resource.result.content', 'Sensitive resource content');
228+
229+
// Other PII attributes should be filtered
202230
expect(result).not.toHaveProperty('mcp.logging.message');
203231
expect(result).not.toHaveProperty('mcp.resource.uri');
204232

233+
// Non-PII attributes should remain
205234
expect(result).toHaveProperty('mcp.method.name', 'tools/call');
206235
expect(result).toHaveProperty('mcp.session.id', 'test-session-123');
207236
});

0 commit comments

Comments
 (0)