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
@@ -299,6 +300,11 @@ def response_hook(span: Span, status: str, response_headers: List):
 
 _logger = getLogger(__name__)
 
+# Global constants for Flask 3.1+ streaming context cleanup
+_IS_FLASK_31_PLUS = hasattr(flask, "__version__") and package_version.parse(
+    flask.__version__
+) >= package_version.parse("3.1.0")
+
 _ENVIRON_STARTTIME_KEY = "opentelemetry-flask.starttime_key"
 _ENVIRON_SPAN_KEY = "opentelemetry-flask.span_key"
 _ENVIRON_ACTIVATION_KEY = "opentelemetry-flask.activation_key"
@@ -408,6 +414,11 @@ def _start_response(status, response_headers, *args, **kwargs):
             return start_response(status, response_headers, *args, **kwargs)
 
         result = wsgi_app(wrapped_app_environ, _start_response)
+
+        # Note: Streaming response context cleanup is now handled in the Flask teardown function
+        # (_wrapped_teardown_request) to ensure proper cleanup following Logfire's recommendations
+        # for OpenTelemetry generator context management
+
         if should_trace:
             duration_s = default_timer() - start
             if duration_histogram_old:
@@ -433,6 +444,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 +549,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 +567,79 @@ 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
+            # Only check for streaming in Flask 3.1+ and Python 3.10+ to avoid interference with older versions
+            is_flask_31_plus = _IS_FLASK_31_PLUS and sys.version_info >= (
+                3,
+                10,
             )
 
-        if flask.request.environ.get(_ENVIRON_TOKEN, None):
-            context.detach(flask.request.environ.get(_ENVIRON_TOKEN))
+            is_streaming = False
+            if is_flask_31_plus:
+                try:
+                    # Additional safety check: verify we're in a Flask request context
+                    if hasattr(flask, "request") and hasattr(
+                        flask.request, "response"
+                    ):
+                        is_streaming = (
+                            hasattr(flask.request, "response")
+                            and flask.request.response
+                            and hasattr(flask.request.response, "stream")
+                            and flask.request.response.stream
+                        )
+                except (RuntimeError, AttributeError):
+                    # Not in a proper Flask request context, don't check for streaming
+                    is_streaming = False
+
+            if is_flask_31_plus and is_streaming:
+                # For Flask 3.1+ streaming responses, ensure OpenTelemetry contexts are cleaned up
+                # This addresses the generator context leak issues documented by Logfire
+                # (open-telemetry/opentelemetry-python#2606)
+                try:
+                    context.detach(token)
+                    if hasattr(activation, "__exit__"):
+                        activation.__exit__(None, None, None)
+
+                    # Mark as cleaned up
+                    flask.request.environ[_ENVIRON_ACTIVATION_KEY] = None
+                    flask.request.environ[_ENVIRON_TOKEN] = None
+
+                    _logger.debug(
+                        "Streaming response context cleanup completed in teardown function"
+                    )
+
+                except (
+                    RuntimeError,
+                    ValueError,
+                    TypeError,
+                    AttributeError,
+                ) as cleanup_exc:
+                    _logger.debug(
+                        "Teardown streaming context cleanup failed: %s",
+                        cleanup_exc,
+                    )
+                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
 

instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import ThreadPoolExecutor, as_completed  # pylint: disable=E0611
 from random import randint
 
 import flask