use nullcontext and add logic to sync_wrapper

Author: alexander-alderman-webb

sentry_sdk/integrations/_wsgi_common.py

@@ -1,4 +1,3 @@
-from contextlib import contextmanager
 import json
 from copy import deepcopy
 
@@ -52,12 +51,6 @@
 )
 
 
-# This noop context manager can be replaced with "from contextlib import nullcontext" when we drop Python 3.6 support
-@contextmanager
-def nullcontext() -> "Iterator[None]":
-    yield
-
-
 def request_body_within_bounds(
     client: "Optional[sentry_sdk.client.BaseClient]", content_length: int
 ) -> bool:

sentry_sdk/integrations/asgi.py

@@ -19,7 +19,6 @@
 )
 from sentry_sdk.integrations._wsgi_common import (
     DEFAULT_HTTP_METHODS_TO_CAPTURE,
-    nullcontext,
 )
 from sentry_sdk.sessions import track_session
 from sentry_sdk.tracing import (
@@ -34,6 +33,7 @@
     logger,
     transaction_from_function,
     _get_installed_modules,
+    nullcontext,
 )
 from sentry_sdk.tracing import Transaction
 

sentry_sdk/integrations/mcp.py

@@ -17,6 +17,7 @@
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.utils import safe_serialize
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.utils import nullcontext
 
 try:
     from mcp.server.lowlevel import Server  # type: ignore[import-not-found]
@@ -80,8 +81,8 @@ def _get_active_http_scopes() -> (
         return None
 
     return (
-        ctx.request.scope["state"].get("sentry_sdk.current_scope"),
         ctx.request.scope["state"].get("sentry_sdk.isolation_scope"),
+        ctx.request.scope["state"].get("sentry_sdk.current_scope"),
     )
 
 
@@ -410,17 +411,28 @@ async def _async_handler_wrapper(
     scopes = _get_active_http_scopes()
 
     if scopes is None:
-        current_scope = None
-        isolation_scope = None
+        isolation_scope_context = nullcontext()
+        current_scope_context = nullcontext()
     else:
-        current_scope, isolation_scope = scopes
+        isolation_scope, current_scope = scopes
+
+        isolation_scope_context = (
+            nullcontext()
+            if isolation_scope is None
+            else sentry_sdk.scope.use_isolation_scope(isolation_scope)
+        )
+        current_scope_context = (
+            nullcontext()
+            if current_scope is None
+            else sentry_sdk.scope.use_scope(current_scope)
+        )
 
     # Get request ID, session ID, and transport from context
     request_id, session_id, mcp_transport = _get_request_context_data()
 
     # Start span and execute
-    with sentry_sdk.scope.use_isolation_scope(isolation_scope):
-        with sentry_sdk.scope.use_scope(current_scope):
+    with isolation_scope_context:
+        with current_scope_context:
             with get_start_span_function()(
                 op=OP.MCP_SERVER,
                 name=span_name,
@@ -490,50 +502,71 @@ def _sync_handler_wrapper(
         result_data_key,
     ) = _prepare_handler_data(handler_type, original_args)
 
-    # Start span and execute
-    with get_start_span_function()(
-        op=OP.MCP_SERVER,
-        name=span_name,
-        origin=MCPIntegration.origin,
-    ) as span:
-        # Get request ID, session ID, and transport from context
-        request_id, session_id, mcp_transport = _get_request_context_data()
-
-        # Set input span data
-        _set_span_input_data(
-            span,
-            handler_name,
-            span_data_key,
-            mcp_method_name,
-            arguments,
-            request_id,
-            session_id,
-            mcp_transport,
+    scopes = _get_active_http_scopes()
+
+    if scopes is None:
+        isolation_scope_context = nullcontext()
+        current_scope_context = nullcontext()
+    else:
+        isolation_scope, current_scope = scopes
+
+        isolation_scope_context = (
+            nullcontext()
+            if isolation_scope is None
+            else sentry_sdk.scope.use_isolation_scope(isolation_scope)
+        )
+        current_scope_context = (
+            nullcontext()
+            if current_scope is None
+            else sentry_sdk.scope.use_scope(current_scope)
         )
 
-        # For resources, extract and set protocol
-        if handler_type == "resource":
-            uri = original_args[0]
-            protocol = None
-            if hasattr(uri, "scheme"):
-                protocol = uri.scheme
-            elif handler_name and "://" in handler_name:
-                protocol = handler_name.split("://")[0]
-            if protocol:
-                span.set_data(SPANDATA.MCP_RESOURCE_PROTOCOL, protocol)
+    # Start span and execute
+    with isolation_scope_context:
+        with current_scope_context:
+            with get_start_span_function()(
+                op=OP.MCP_SERVER,
+                name=span_name,
+                origin=MCPIntegration.origin,
+            ) as span:
+                # Get request ID, session ID, and transport from context
+                request_id, session_id, mcp_transport = _get_request_context_data()
 
-        try:
-            # Execute the sync handler
-            result = func(*original_args)
-        except Exception as e:
-            # Set error flag for tools
-            if handler_type == "tool":
-                span.set_data(SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True)
-            sentry_sdk.capture_exception(e)
-            raise
-
-        _set_span_output_data(span, result, result_data_key, handler_type)
-        return result
+                # Set input span data
+                _set_span_input_data(
+                    span,
+                    handler_name,
+                    span_data_key,
+                    mcp_method_name,
+                    arguments,
+                    request_id,
+                    session_id,
+                    mcp_transport,
+                )
+
+                # For resources, extract and set protocol
+                if handler_type == "resource":
+                    uri = original_args[0]
+                    protocol = None
+                    if hasattr(uri, "scheme"):
+                        protocol = uri.scheme
+                    elif handler_name and "://" in handler_name:
+                        protocol = handler_name.split("://")[0]
+                    if protocol:
+                        span.set_data(SPANDATA.MCP_RESOURCE_PROTOCOL, protocol)
+
+                try:
+                    # Execute the sync handler
+                    result = func(*original_args)
+                except Exception as e:
+                    # Set error flag for tools
+                    if handler_type == "tool":
+                        span.set_data(SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True)
+                    sentry_sdk.capture_exception(e)
+                    raise
+
+                _set_span_output_data(span, result, result_data_key, handler_type)
+                return result
 
 
 def _create_instrumented_handler(

sentry_sdk/integrations/wsgi.py

@@ -1,4 +1,5 @@
 import sys
+import functools
 from functools import partial
 from typing import TYPE_CHECKING
 
@@ -9,7 +10,6 @@
 from sentry_sdk.integrations._wsgi_common import (
     DEFAULT_HTTP_METHODS_TO_CAPTURE,
     _filter_headers,
-    nullcontext,
 )
 from sentry_sdk.scope import should_send_default_pii, use_isolation_scope
 from sentry_sdk.sessions import track_session
@@ -19,6 +19,7 @@
     capture_internal_exceptions,
     event_from_exception,
     reraise,
+    nullcontext,
 )
 
 if TYPE_CHECKING:
@@ -70,6 +71,7 @@ class SentryWsgiMiddleware:
         "use_x_forwarded_for",
         "span_origin",
         "http_methods_to_capture",
+        "_wsgi_file_wrapper",
     )
 
     def __init__(
@@ -85,11 +87,19 @@ def __init__(
         self.http_methods_to_capture = http_methods_to_capture
 
     def __call__(
-        self, environ: "Dict[str, str]", start_response: "Callable[..., Any]"
+        self, environ: "Dict[str, Any]", start_response: "Callable[..., Any]"
     ) -> "_ScopedResponse":
         if _wsgi_middleware_applied.get(False):
             return self.app(environ, start_response)
 
+        old_wsgi_file_wrapper = environ["wsgi.file_wrapper"]
+
+        def _patch_environ_wsgi_file_wrapper(fileobj, block_size=8192):
+            self._wsgi_file_wrapper = old_wsgi_file_wrapper(fileobj, block_size)
+            return self._wsgi_file_wrapper
+
+        environ["wsgi.file_wrapper"] = _patch_environ_wsgi_file_wrapper
+
         _wsgi_middleware_applied.set(True)
         try:
             with sentry_sdk.isolation_scope() as scope:
@@ -135,6 +145,9 @@ def __call__(
         finally:
             _wsgi_middleware_applied.set(False)
 
+        if response is self._wsgi_file_wrapper:
+            return response
+
         return _ScopedResponse(scope, response)
 
 
@@ -215,6 +228,14 @@ def _capture_exception() -> "ExcInfo":
     return exc_info
 
 
+class _ScopedFileWrapper:
+    def __init__(self, scope: "sentry_sdk.scope.Scope", fileno):
+        pass
+
+    def fileno():
+        pass
+
+
 class _ScopedResponse:
     """
     Users a separate scope for each response chunk.

sentry_sdk/scope.py

@@ -100,7 +100,6 @@
 
     F = TypeVar("F", bound=Callable[..., Any])
     T = TypeVar("T")
-    S = TypeVar("S", bound=Optional["Scope"])
 
 
 # Holds data that will be added to **all** events sent by this process.
@@ -1787,7 +1786,7 @@ def new_scope() -> "Generator[Scope, None, None]":
 
 
 @contextmanager
-def use_scope(scope: "S") -> "Generator[S, None, None]":
+def use_scope(scope: "Scope") -> "Generator[Scope, None, None]":
     """
     .. versionadded:: 2.0.0
 
@@ -1809,10 +1808,6 @@ def use_scope(scope: "S") -> "Generator[S, None, None]":
         sentry_sdk.capture_message("hello, again") # will NOT include `color` tag.
 
     """
-    if scope is None:
-        yield scope
-        return
-
     # set given scope as current scope
     token = _current_scope.set(scope)
 
@@ -1876,7 +1871,7 @@ def isolation_scope() -> "Generator[Scope, None, None]":
 
 
 @contextmanager
-def use_isolation_scope(isolation_scope: "S") -> "Generator[S, None, None]":
+def use_isolation_scope(isolation_scope: "Scope") -> "Generator[Scope, None, None]":
     """
     .. versionadded:: 2.0.0
 
@@ -1897,10 +1892,6 @@ def use_isolation_scope(isolation_scope: "S") -> "Generator[S, None, None]":
         sentry_sdk.capture_message("hello, again") # will NOT include `color` tag.
 
     """
-    if isolation_scope is None:
-        yield isolation_scope
-        return
-
     # fork current scope
     current_scope = Scope.get_current_scope()
     forked_current_scope = current_scope.fork()

sentry_sdk/utils.py

@@ -16,6 +16,7 @@
 from functools import partial, partialmethod, wraps
 from numbers import Real
 from urllib.parse import parse_qs, unquote, urlencode, urlsplit, urlunsplit
+from contextlib import contextmanager
 
 try:
     # Python 3.11
@@ -2078,3 +2079,9 @@ def serialize_attribute(val: "AttributeValue") -> "SerializedAttributeValue":
     # Coerce to string if we don't know what to do with the value. This should
     # never happen as we pre-format early in format_attribute, but let's be safe.
     return {"value": safe_repr(val), "type": "string"}
+
+
+# This noop context manager can be replaced with "from contextlib import nullcontext" when we drop Python 3.6 support
+@contextmanager
+def nullcontext() -> "Iterator[None]":
+    yield