lib/with_popup/rack.rb in with_popup-0.0.6 vs lib/with_popup/rack.rb in with_popup-0.0.7

@@ -5,11 +5,11 @@
 
   def call(env)
     @env = env
     @status, @headers, @response = @app.call(env)
     return [@status, @headers, @response] unless need_injection?
-    body = inject_js(@response.body)
+    body = inject_js_to_the_response
     clear_param!
     [@status, @headers, [body]]
   end
 
   private
@@ -32,10 +32,11 @@
 
   def clear_param!
     session.delete :_with_popup
   end
 
-  def inject_js(body)
+  def inject_js_to_the_response
+    body = @response.try(:body) || @response.first
     js = popup_param == :close ? close_popup_js : reload_popup_js
     body.sub('</body>', js + '</body>')
   end
 
   def close_popup_js