ClearEventReporterContext handles case where rack.response_finished is not supported

Author: adrianna-chang-shopify

actionpack/lib/action_dispatch/middleware/clear_event_reporter_context.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 module ActionDispatch
+  # TODO: This should go away if we ship https://github.com/rails/rails/pull/55425
+  #
   # Middleware that sets up a callback on rack.response_finished to clear
   # the EventReporter context when the response is finished. This ensures that
   # context is cleared as late as possible in the request lifecycle.
@@ -10,11 +12,18 @@ def initialize(app)
     end
 
     def call(env)
+      if response_finished = env["rack.response_finished"]
+        response_finished << -> do
+          ActiveSupport.event_reporter.clear_context
+        end
+      end
+
       response = @app.call(env)
 
-      env["rack.response_finished"] ||= []
-      env["rack.response_finished"] << -> do
-        ActiveSupport.event_reporter.clear_context
+      unless response_finished
+        response << ::Rack::BodyProxy.new(response.pop) do
+          ActiveSupport.event_reporter.clear_context
+        end
       end
 
       response

actionpack/test/dispatch/clear_event_reporter_context_test.rb

@@ -5,15 +5,14 @@
 class ClearEventReporterContextTest < ActiveSupport::TestCase
   def setup
     @app = ->(env) { [200, {}, ["Hello"]] }
-    @middleware = ActionDispatch::ClearEventReporterContext.new(@app)
     @reporter = ActiveSupport.event_reporter
   end
 
   test "clears event reporter context in response finished callback" do
     @reporter.set_context(shop_id: 123)
 
-    env = {}
-    @middleware.call(env)
+    env = Rack::MockRequest.env_for("", { "rack.response_finished" => [] })
+    middleware(@app).call(env)
 
     assert env["rack.response_finished"]
     assert_equal 1, env["rack.response_finished"].length
@@ -23,17 +22,46 @@ def setup
     assert_equal({}, @reporter.context)
   end
 
-  test "clears event reporter context when exception is raised" do
+  test "clears event reporter context via Rack::BodyProxy when rack.response_finished is not supported" do
     @reporter.set_context(shop_id: 123)
 
+    env = Rack::MockRequest.env_for("", {})
+    response = middleware(@app).call(env)
+
+    assert_equal({ shop_id: 123 }, @reporter.context)
+
+    body = response[2]
+    body.close
+
+    assert_equal({}, @reporter.context)
+  end
+
+  test "clears event reporter context via rack.response_finished when exception is raised" do
+    @reporter.set_context(shop_id: 123)
+
+    env = Rack::MockRequest.env_for("", { "rack.response_finished" => [] })
     exception_app = ->(env) { raise StandardError, "Test exception" }
-    exception_middleware = ActionDispatch::ClearEventReporterContext.new(exception_app)
+    assert_raises(StandardError) do
+      middleware(exception_app).call(env)
+    end
 
-    env = {}
+    assert_equal({}, @reporter.context)
+  end
+
+  test "clears event reporter context when exception is raised and rack.response_finished is not supported" do
+    @reporter.set_context(shop_id: 123)
+
+    exception_app = ->(env) { raise StandardError, "Test exception" }
+    env = Rack::MockRequest.env_for("", {})  # Add missing env variable
     assert_raises(StandardError) do
-      exception_middleware.call(env)
+      middleware(exception_app).call(env)
     end
 
     assert_equal({}, @reporter.context)
   end
+
+  private
+    def middleware(inner_app)
+      Rack::Lint.new(ActionDispatch::ClearEventReporterContext.new(Rack::Lint.new(inner_app)))
+    end
 end