Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 71 additions & 23 deletions sentry_sdk/integrations/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,12 @@
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing_utils import has_span_streaming_enabled
from sentry_sdk.utils import nullcontext, package_version, safe_serialize
from sentry_sdk.utils import (
has_data_collection_enabled,
nullcontext,
package_version,
safe_serialize,
)

try:
from mcp.server.lowlevel import Server
Expand Down Expand Up @@ -322,13 +327,21 @@ async def _tool_handler_wrapper(
original_kwargs: Original keyword arguments passed to the handler
self: Optional instance for bound methods
"""
client = sentry_sdk.get_client()

if original_kwargs is None:
original_kwargs = {}

handler_name, arguments = _extract_handler_data_from_args(
"tool", original_args, original_kwargs
)

if has_data_collection_enabled(client.options):
if not client.options["data_collection"]["gen_ai"]["inputs"]:
# Arguments can contain sensitive data and shouldn't be added to the span
# if the user has opted out via this config.
arguments = {}

ctx = None
try:
ctx = request_ctx.get()
Expand Down Expand Up @@ -388,14 +401,17 @@ async def _tool_handler_wrapper(
return result

# Get integration to check PII settings
integration = sentry_sdk.get_client().get_integration(MCPIntegration)
integration = client.get_integration(MCPIntegration)
if integration is None:
return result

# Check if we should include sensitive data
should_include_data = (
should_send_default_pii() and integration.include_prompts
)
should_include_data = False
if has_data_collection_enabled(client.options):
if client.options["data_collection"]["gen_ai"]["outputs"]:
should_include_data = True
elif should_send_default_pii() and integration.include_prompts:
should_include_data = True

extracted = _extract_tool_result_content(result)
if extracted is not None and should_include_data:
Expand All @@ -422,15 +438,22 @@ async def _instrument_v2_tool_call(
if ctx.params is None or ctx.params.get("name") is None:
return await call_next(ctx)

client = sentry_sdk.get_client()
handler_name = ctx.params["name"]
arguments = ctx.params.get("arguments")
if arguments is None:
arguments = {}

if has_data_collection_enabled(client.options):
if not client.options["data_collection"]["gen_ai"]["inputs"]:
# Arguments can contain sensitive data and shouldn't be added to the span
# if the user has opted out via this config.
arguments = {}

# Get request ID, session ID, and transport from context
request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx)

span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
span_streaming = has_span_streaming_enabled(client.options)

# Start span and execute
with _active_http_scopes(ctx=ctx):
Expand Down Expand Up @@ -474,22 +497,25 @@ async def _instrument_v2_tool_call(
return result

# Get integration to check PII settings
integration = sentry_sdk.get_client().get_integration(MCPIntegration)
integration = client.get_integration(MCPIntegration)
if integration is None:
return result

# Check if we should include sensitive data
should_include_data = (
should_send_default_pii() and integration.include_prompts
)
should_include_result_data = False
if has_data_collection_enabled(client.options):
if client.options["data_collection"]["gen_ai"]["outputs"]:
should_include_result_data = True
elif should_send_default_pii() and integration.include_prompts:
should_include_result_data = True

result_content = result
if "structuredContent" in result:
result_content = result["structuredContent"]
elif isinstance(result.get("content"), list):
result_content = _extract_text_from_content_blocks(result["content"])

if result_content is not None and should_include_data:
if result_content is not None and should_include_result_data:
_set_span_data_attribute(
span,
SPANDATA.MCP_TOOL_RESULT_CONTENT,
Expand Down Expand Up @@ -526,10 +552,17 @@ async def _prompt_handler_wrapper(
if original_kwargs is None:
original_kwargs = {}

client = sentry_sdk.get_client()
handler_name, arguments = _extract_handler_data_from_args(
"prompt", original_args, original_kwargs
)

if has_data_collection_enabled(client.options):
if not client.options["data_collection"]["gen_ai"]["inputs"]:
# Arguments can contain sensitive data and shouldn't be added to the span
# if the user has opted out via this config.
arguments = {}

ctx = None
try:
ctx = request_ctx.get()
Expand All @@ -539,7 +572,7 @@ async def _prompt_handler_wrapper(
# Get request ID, session ID, and transport from context
request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx)

span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
span_streaming = has_span_streaming_enabled(client.options)

# Start span and execute
with _active_http_scopes(ctx=ctx):
Expand Down Expand Up @@ -589,14 +622,17 @@ async def _prompt_handler_wrapper(
return result

# Get integration to check PII settings
integration = sentry_sdk.get_client().get_integration(MCPIntegration)
integration = client.get_integration(MCPIntegration)
if integration is None:
return result

# Check if we should include sensitive data
should_include_data = (
should_send_default_pii() and integration.include_prompts
)
should_include_result_data = False
if has_data_collection_enabled(client.options):
if client.options["data_collection"]["gen_ai"]["inputs"]:
should_include_result_data = True
elif should_send_default_pii() and integration.include_prompts:
should_include_result_data = True

# For prompts, count messages and set role/content only for single-message prompts
try:
Expand All @@ -619,7 +655,7 @@ async def _prompt_handler_wrapper(
)

# Only set role and content for single-message prompts if PII is allowed
if message_count == 1 and should_include_data and messages:
if message_count == 1 and should_include_result_data and messages:
first_message = messages[0]
# Extract role
role = None
Expand Down Expand Up @@ -675,15 +711,24 @@ async def _instrument_v2_prompt_get(
if ctx.params is None or ctx.params.get("name") is None:
return await call_next(ctx)

client = sentry_sdk.get_client()
handler_name = ctx.params["name"]

arguments = ctx.params.get("arguments")

if arguments is None:
arguments = {}

if has_data_collection_enabled(client.options):
if not client.options["data_collection"]["gen_ai"]["inputs"]:
# Arguments can contain sensitive data and shouldn't be added to the span
# if the user has opted out via this config.
arguments = {}

# Get request ID, session ID, and transport from context
request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx)

span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
span_streaming = has_span_streaming_enabled(client.options)

# Start span and execute
with _active_http_scopes(ctx=ctx):
Expand Down Expand Up @@ -726,14 +771,17 @@ async def _instrument_v2_prompt_get(
return result

# Get integration to check PII settings
integration = sentry_sdk.get_client().get_integration(MCPIntegration)
integration = client.get_integration(MCPIntegration)
if integration is None:
return result

# Check if we should include sensitive data
should_include_data = (
should_send_default_pii() and integration.include_prompts
)
should_include_result_data = False
if has_data_collection_enabled(client.options):
if client.options["data_collection"]["gen_ai"]["inputs"]:
should_include_result_data = True
elif should_send_default_pii() and integration.include_prompts:
should_include_result_data = True

# For prompts, count messages and set role/content only for single-message prompts
try:
Expand All @@ -751,7 +799,7 @@ async def _instrument_v2_prompt_get(
)

# Only set role and content for single-message prompts if PII is allowed
if message_count == 1 and should_include_data and messages:
if message_count == 1 and should_include_result_data and messages:
first_message = messages[0]
# Extract role
role = None
Expand Down
Loading
Loading