5 files changed, 57 insertions, 12 deletions
diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index fe63cf2b98..41b997a755 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -125,9 +125,11 @@ module ActionController
       end
 
       def each
+        @response.sending!
         while str = @buf.pop
           yield str
         end
+        @response.sent!
       end
 
       def close
@@ -179,12 +181,16 @@ module ActionController
 
       private
 
-      def finalize_response
+      def before_committed
         super
         jar = request.cookie_jar
         # The response can be committed multiple times
         jar.write self unless committed?
-        jar.commit!
+      end
+
+      def before_sending
+        super
+        request.cookie_jar.commit!
         headers.freeze
       end
 
diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb
index 009f83861d..e9166d8747 100644
--- a/actionpack/lib/action_controller/test_case.rb
+++ b/actionpack/lib/action_controller/test_case.rb
@@ -595,7 +595,7 @@ module ActionController
         @controller.process(name)
 
         if cookies = @request.env['action_dispatch.cookies']
-          unless cookies.committed?
+          unless @response.committed?
             cookies.write(@response)
           end
         end
diff --git a/actionpack/lib/action_dispatch/http/response.rb b/actionpack/lib/action_dispatch/http/response.rb
index b5454f519f..3d27ff2b24 100644
--- a/actionpack/lib/action_dispatch/http/response.rb
+++ b/actionpack/lib/action_dispatch/http/response.rb
@@ -91,7 +91,10 @@ module ActionDispatch # :nodoc:
       end
 
       def each(&block)
-        @buf.each(&block)
+        @response.sending!
+        x = @buf.each(&block)
+        @response.sent!
+        x
       end
 
       def close
@@ -118,6 +121,8 @@ module ActionDispatch # :nodoc:
       @blank        = false
       @cv           = new_cond
       @committed    = false
+      @sending      = false
+      @sent         = false
       @content_type = nil
       @charset      = nil
 
@@ -138,18 +143,37 @@ module ActionDispatch # :nodoc:
       end
     end
 
+    def await_sent
+      synchronize { @cv.wait_until { @sent } }
+    end
+
     def commit!
       synchronize do
-        finalize_response
+        before_committed
         @committed = true
         @cv.broadcast
       end
     end
 
-    def committed?
-      @committed
+    def sending!
+      synchronize do
+        before_sending
+        @sending = true
+        @cv.broadcast
+      end
     end
 
+    def sent!
+      synchronize do
+        @sent = true
+        @cv.broadcast
+      end
+    end
+
+    def sending?;   synchronize { @sending };   end
+    def committed?; synchronize { @committed }; end
+    def sent?;      synchronize { @sent };      end
+
     # Sets the HTTP status code.
     def status=(status)
       @status = Rack::Utils.status_code(status)
@@ -274,7 +298,10 @@ module ActionDispatch # :nodoc:
 
   private
 
-    def finalize_response
+    def before_committed
+    end
+
+    def before_sending
     end
 
     def merge_default_headers(original, default)
diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index cc232f6b6e..2fd5c930ba 100644
--- a/actionpack/test/controller/live_stream_test.rb
+++ b/actionpack/test/controller/live_stream_test.rb
@@ -44,7 +44,7 @@ module ActionController
     tests SSETestController
 
     def wait_for_response_stream_close
-      response.stream.await_close
+      response.body
     end
 
     def test_basic_sse
@@ -186,8 +186,8 @@ module ActionController
     tests TestController
 
     def assert_stream_closed
-      response.stream.await_close
       assert response.stream.closed?, 'stream should be closed'
+      assert response.sent?, 'stream should be sent'
     end
 
     def capture_log_output
@@ -229,6 +229,7 @@ module ActionController
       @controller.response = @response
 
       t = Thread.new(@response) { |resp|
+        resp.await_commit
         resp.stream.each do |part|
           assert_equal parts.shift, part
           ol = @controller.latch
@@ -268,6 +269,7 @@ module ActionController
       assert_raises(ActionView::MissingTemplate) do
         get :exception_in_view
       end
+      assert response.body
       assert_stream_closed
     end
 
diff --git a/actionpack/test/dispatch/live_response_test.rb b/actionpack/test/dispatch/live_response_test.rb
index 336fdd569b..512f3a8a7a 100644
--- a/actionpack/test/dispatch/live_response_test.rb
+++ b/actionpack/test/dispatch/live_response_test.rb
@@ -35,6 +35,7 @@ module ActionController
           @response.stream.close
         }
 
+        @response.await_commit
         @response.each do |part|
           assert_equal 'foo', part
           latch.release
@@ -59,21 +60,30 @@ module ActionController
         assert_nil @response.headers['Content-Length']
       end
 
-      def test_headers_cannot_be_written_after_write
+      def test_headers_cannot_be_written_after_webserver_reads
         @response.stream.write 'omg'
+        latch = ActiveSupport::Concurrency::Latch.new
 
+        t = Thread.new {
+          @response.stream.each do |chunk|
+            latch.release
+          end
+        }
+
+        latch.await
         assert @response.headers.frozen?
         e = assert_raises(ActionDispatch::IllegalStateError) do
           @response.headers['Content-Length'] = "zomg"
         end
 
         assert_equal 'header already sent', e.message
+        @response.stream.close
+        t.join
       end
 
       def test_headers_cannot_be_written_after_close
         @response.stream.close
 
-        assert @response.headers.frozen?
         e = assert_raises(ActionDispatch::IllegalStateError) do
           @response.headers['Content-Length'] = "zomg"
         end