chore: compatiable flask 3.1+

Author: fatelei

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

@@ -333,6 +333,49 @@ 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.
+    """
+    activation = environ.get(_ENVIRON_ACTIVATION_KEY)
+    token = environ.get(_ENVIRON_TOKEN)
+
+    if not activation or not token:
+        return
+
+    try:
+        # Check if we have a Flask request context and if this might be a streaming response
+        # For Flask 3.1+, we need to be more proactive about context cleanup
+        # because streaming responses can cause token reuse issues
+
+        # Mark that we've handled the cleanup to prevent double cleanup in teardown
+        environ[_ENVIRON_ACTIVATION_KEY] = None
+        environ[_ENVIRON_TOKEN] = None
+
+        # Clean up the context token safely
+        if token:
+            try:
+                context.detach(token)
+            except RuntimeError as e:
+                # 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", e)
+
+        # Clean up the activation
+        if hasattr(activation, "__exit__"):
+            try:
+                activation.__exit__(None, None, None)
+            except Exception as e:
+                _logger.debug("Error during activation cleanup: %s", e)
+
+    except Exception:
+        # Log the error but don't raise it to avoid breaking the response
+        _logger.debug("Error during streaming context cleanup", exc_info=True)
+
+
 def _rewrapped_app(
     wsgi_app,
     active_requests_counter,
@@ -408,6 +451,13 @@ 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
+        if should_trace and package_version.parse(
+            flask_version
+        ) >= package_version.parse("3.1.0"):
+            _ensure_streaming_context_cleanup(wrapped_app_environ)
+
         if should_trace:
             duration_s = default_timer() - start
             if duration_histogram_old:
@@ -433,6 +483,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 +588,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 +606,42 @@ 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)
+
+        try:
+            # For Flask 3.1+, check if this is a streaming response that might
+            # have already been cleaned up to prevent double cleanup
+            is_streaming = (
+                hasattr(flask.request, "response")
+                and flask.request.response
+                and hasattr(flask.request.response, "stream")
+                and flask.request.response.stream
             )
 
-        if flask.request.environ.get(_ENVIRON_TOKEN, None):
-            context.detach(flask.request.environ.get(_ENVIRON_TOKEN))
+            if (
+                package_version.parse(flask_version)
+                >= package_version.parse("3.1.0")
+                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 Exception:
+            # Log the error but don't raise it to avoid breaking the request handling
+            _logger.debug("Error during request teardown", exc_info=True)
 
     return _teardown_request