lib/goliath/rack/default_response_format.rb in goliath-0.9.1 vs lib/goliath/rack/default_response_format.rb in goliath-0.9.2

@@ -1,23 +1,11 @@
 module Goliath
   module Rack
     class DefaultResponseFormat
-      def initialize(app)
-        @app = app
-      end
+      include Goliath::Rack::AsyncMiddleware
 
-      def call(env)
-        async_cb = env['async.callback']
-        env['async.callback'] = Proc.new do |status, headers, body|
-          async_cb.call(post_process(status, headers, body))
-        end
-
-        status, headers, body = @app.call(env)
-        post_process(status, headers, body)
-      end
-
-      def post_process(status, headers, body)
+      def post_process(env, status, headers, body)
         return [status, headers, body] if body.respond_to?(:to_ary)
 
         new_body = []
         if body.respond_to?(:each)
           body.each { |chunk| new_body << chunk }
@@ -28,6 +16,6 @@
 
         [status, headers, new_body.flatten]
       end
     end
   end
-end
\ No newline at end of file
+end