chore: compatiable flask 3.1+

Author: fatelei

instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py

@@ -254,6 +254,7 @@ def response_hook(span: Span, status: str, response_headers: List):
 ---
 """
 
+import sys
 import weakref
 from logging import getLogger
 from time import time_ns
@@ -333,6 +334,66 @@ def get_default_span_name():
     return span_name
 
 
+def _ensure_streaming_context_cleanup(environ):
+    """
+    Ensure proper context cleanup for streaming responses in Flask 3.1+.
+
+    This function checks if the response is a streaming response and ensures
+    that context tokens are properly cleaned up to prevent token reuse issues.
+    Only applies to Flask 3.1+ and Python 3.10+ for compatibility reasons.
+    """
+    # Double-check Flask version - this should only run in Flask 3.1+
+    if package_version.parse(flask_version) < package_version.parse("3.1.0"):
+        return
+
+    # Only enable streaming context cleanup for Python 3.10+ to avoid compatibility issues
+    # with older Python versions that have different context management behavior
+    if sys.version_info < (3, 10):
+        return
+
+    activation = environ.get(_ENVIRON_ACTIVATION_KEY)
+    token = environ.get(_ENVIRON_TOKEN)
+
+    if not activation or not token:
+        return
+
+    # Additional safety check - only proceed if we haven't already cleaned up
+    if (
+        environ.get(_ENVIRON_ACTIVATION_KEY) is None
+        or environ.get(_ENVIRON_TOKEN) is None
+    ):
+        return
+
+    try:
+        # Clean up the context token safely
+        if token:
+            try:
+                context.detach(token)
+            except RuntimeError as exc:
+                # Token has already been used - this can happen in Flask 3.1+
+                # with streaming responses, so we just log and continue
+                _logger.debug("Token already detached, continuing: %s", exc)
+                # If detach failed, don't proceed with activation cleanup
+                return
+
+        # Clean up the activation
+        if hasattr(activation, "__exit__"):
+            try:
+                activation.__exit__(None, None, None)
+            except (RuntimeError, AttributeError) as exc:
+                _logger.debug("Error during activation cleanup: %s", exc)
+
+        # Mark that we've handled the cleanup to prevent double cleanup in teardown
+        environ[_ENVIRON_ACTIVATION_KEY] = None
+        environ[_ENVIRON_TOKEN] = None
+
+    except (RuntimeError, ValueError, TypeError) as exc:
+        # Log the error but don't raise it to avoid breaking the response
+        _logger.debug(
+            "Error during streaming context cleanup: %s", exc, exc_info=True
+        )
+
+
 def _rewrapped_app(
     wsgi_app,
     active_requests_counter,
@@ -408,6 +469,37 @@ def _start_response(status, response_headers, *args, **kwargs):
             return start_response(status, response_headers, *args, **kwargs)
 
         result = wsgi_app(wrapped_app_environ, _start_response)
+
+        # For Flask 3.1+, check if we need to handle streaming response context cleanup
+        # Only run this logic in Flask 3.1+ and Python 3.10+ to avoid any interference with older versions
+        if (
+            should_trace
+            and package_version.parse(flask_version)
+            >= package_version.parse("3.1.0")
+            and sys.version_info >= (3, 10)
+        ):
+            # Only call streaming context cleanup for actual streaming responses
+            is_streaming = (
+                hasattr(flask.request, "response")
+                and flask.request.response
+                and hasattr(flask.request.response, "stream")
+                and flask.request.response.stream
+            )
+            if is_streaming:
+                try:
+                    _ensure_streaming_context_cleanup(wrapped_app_environ)
+                except (
+                    RuntimeError,
+                    ValueError,
+                    TypeError,
+                    AttributeError,
+                ) as exc:
+                    # Ensure our Flask 3.1+ logic never interferes with normal request processing
+                    _logger.debug(
+                        "Flask 3.1+ streaming context cleanup failed, continuing: %s",
+                        exc,
+                    )
+
         if should_trace:
             duration_s = default_timer() - start
             if duration_histogram_old:
@@ -433,6 +525,7 @@ def _start_response(status, response_headers, *args, **kwargs):
                 duration_histogram_new.record(
                     max(duration_s, 0), duration_attrs_new
                 )
+
         active_requests_counter.add(-1, active_requests_count_attrs)
         return result
 
@@ -537,6 +630,7 @@ def _teardown_request(exc):
             return
 
         activation = flask.request.environ.get(_ENVIRON_ACTIVATION_KEY)
+        token = flask.request.environ.get(_ENVIRON_TOKEN)
 
         original_reqctx_ref = flask.request.environ.get(
             _ENVIRON_REQCTX_REF_KEY
@@ -554,15 +648,49 @@ def _teardown_request(exc):
             # like any decorated with `flask.copy_current_request_context`.
 
             return
-        if exc is None:
-            activation.__exit__(None, None, None)
-        else:
-            activation.__exit__(
-                type(exc), exc, getattr(exc, "__traceback__", None)
-            )
 
-        if flask.request.environ.get(_ENVIRON_TOKEN, None):
-            context.detach(flask.request.environ.get(_ENVIRON_TOKEN))
+        try:
+            # For Flask 3.1+, check if this is a streaming response that might
+            # have already been cleaned up to prevent double cleanup
+            # Only check for streaming in Flask 3.1+ and Python 3.10+ to avoid interference with older versions
+            is_flask_31_plus = package_version.parse(
+                flask_version
+            ) >= package_version.parse("3.1.0") and sys.version_info >= (3, 10)
+
+            is_streaming = False
+            if is_flask_31_plus:
+                is_streaming = (
+                    hasattr(flask.request, "response")
+                    and flask.request.response
+                    and hasattr(flask.request.response, "stream")
+                    and flask.request.response.stream
+                )
+
+            if is_flask_31_plus and is_streaming:
+                # For streaming responses in Flask 3.1+, the context might have been
+                # cleaned up already in _ensure_streaming_context_cleanup
+                # Mark the activation and token as None to prevent double cleanup
+                flask.request.environ[_ENVIRON_ACTIVATION_KEY] = None
+                flask.request.environ[_ENVIRON_TOKEN] = None
+                return
+
+            if exc is None:
+                activation.__exit__(None, None, None)
+            else:
+                activation.__exit__(
+                    type(exc), exc, getattr(exc, "__traceback__", None)
+                )
+
+            if token:
+                context.detach(token)
+
+        except (RuntimeError, AttributeError, ValueError) as teardown_exc:
+            # Log the error but don't raise it to avoid breaking the request handling
+            _logger.debug(
+                "Error during request teardown: %s",
+                teardown_exc,
+                exc_info=True,
+            )
 
     return _teardown_request