template_streaming 0.0.11 → 0.1.0

data/CHANGELOG

@@ -1,3 +1,12 @@
+== 0.1.0 2011-05-16
+
+ * Replace prelayouts with :stream => true option to #render.
+ * Controller-level #stream to stream selected actions.
+ * Inject errors during rendering to end of body.
+ * Autoflushing.
+ * Fix page and action caching for streamed responses.
+ * New Relic support.
+
 == 0.0.11 2010-05-10
 
  * Ensure authenticity token is initialized properly.

data/README.markdown

@@ -1,6 +1,6 @@
 # Template Streaming
 
-Rails plugin which enables progressive rendering for templates.
+Progressive rendering for Rails.
 
 ## Background
 
@@ -8,33 +8,45 @@ A typical Rails client-side profile looks something like this:
 
 ![Typical Rails Profile][slow-profile]
 
-In almost all cases, this is highly suboptimal, as many resources, such as
-external stylesheets, are static and could be loaded by the client while it's
-waiting for the server response.
+This is highly suboptimal. Many resources, such as external stylesheets, are
+completely static and could be loaded by the client while it's waiting for the
+server response.
 
-The trick is to output the response *progressively*--flushing the stylesheet
-link tags out to the client before it has rendered the rest of the
-page. Depending on how other external resources such javascripts and images are
-used, they too may be flushed out early, significantly reducing the time for the
-page to become interactive.
+The trick is to *stream* the response--flushing the markup for the static
+resources to the client before it has rendered the rest of the page. In
+addition to being able to render styles and images earlier, the browser can
+download javascripts, making the page responsive to input events sooner.
 
-The problem is Rails has never been geared to allow this. Most Rails
-applications use layouts, which require rendering the content of the page before
-the layout. Since the global stylesheet tag is usually in the layout, we can't
-simply flush the rendering buffer from a helper method.
+The main barrier to this in Rails is that layouts are rendered before the
+content of the page. The control flow must thus be altered to render the page in
+the order the client needs to receive it - layout first.
 
-Until now.
-
-Template Streaming circumvents the template rendering order by introducing
-*prelayouts*. A prelayout wraps a layout, and is rendered *before* the layout
-and its content. By using the provided `flush` helper prior to yielding in the
-prelayout, one can now output content early in the rendering process, giving
-profiles that look more like:
+With streaming, your profiles can look more like this:
 
 ![Progressive Rendering Profile][fast-profile]
 
-[slow-profile]: http://github.com/oggy/template_streaming/raw/master/doc/slow-profile.png
-[fast-profile]: http://github.com/oggy/template_streaming/raw/master/doc/fast-profile.png
+[slow-profile]: https://github.com/oggy/template_streaming/raw/master/doc/slow-profile.png
+[fast-profile]: https://github.com/oggy/template_streaming/raw/master/doc/fast-profile.png
+
+## How
+
+Just add the `template_streaming` gem to your application, and add a `stream`
+call for the actions you'd like to stream. For example, to stream just the
+`index` action of your `HomeController`, it would look like this:
+
+    class HomeController
+      stream :only => :index
+
+      def index
+        ...
+      end
+    end
+
+To stream everything, just add `stream` to your `ApplicationController`.
+
+Now you may pepper `flush` calls strategically throughout your views to force a
+flush, such as just after the stylesheet and javascript tags. `flush` may occur
+in both templates and their layouts.
 
 ## API
 
@@ -48,9 +60,12 @@ This has several implications:
  * Anything that needs to inspect or modify the body should be moved to a
    middleware.
  * Modifications to cookies (this includes the flash and session if using the
-   cookie store!) should not be made in the view.
- * An exception during rendering cannot simply replace the body with a
-   stacktrace or 500 page. (Solution to come.)
+   cookie store) must not be made in the view. In fact, these objects will be
+   frozen when streaming.
+ * An exception during rendering cannot result in a 500 response, as the headers
+   will have already been sent. Instead, the innermost partial which contains an
+   error will simply render nothing, and error information is injected into the
+   foot of the page in development mode.
 
 ### Helpers
 
@@ -58,32 +73,76 @@ This has several implications:
     client immediately.
  * `push(data)` - send the given data to the client immediately.
 
-These can only do their job if the underlying web server supports progressive
-rendering via Rack. This has been tested successfully with [Mongrel][mongrel]
-and [Passenger][passenger]. [Thin][thin] is only supported if the [Event Machine
-Flush][event-machine-flush] gem is installed. WEBrick does not support
-progressive rendering. [Please send me][contact] reports of success with other
+## Support
+
+Template Streaming currently only supports Rails 2.3.11. Rails 3.0 support is
+planned in the near future. Rails 3.1 will ship with support for streaming. This
+gem will be updated to meet the API of Rails 3.1 as it evolves, to help you
+migrate.
+
+Streaming also requires a web server that does not buffer Rack responses. It has
+been tested **successfully** with [Passenger][passenger], [Unicorn][unicorn],
+and [Mongrel][mongrel]. Note that Unicorn requires the `:tcp_nopush => false`
+configuration option. [Thin][thin] is only supported if the
+[Event Machine Flush][event-machine-flush] gem is installed. WEBrick does
+**not** support streaming. [Please send me][contact] your experiences with other
 web servers!
 
-[mongrel]: http://github.com/fauna/mongrel
 [passenger]: http://www.modrails.com
-[thin]: http://github.com/macournoyer/thin
-[event-machine-flush]: http://github.com/oggy/event_machine_flush
+[unicorn]: http://unicorn.bogomips.org/
+[mongrel]: https://github.com/fauna/mongrel
+[thin]: https://github.com/macournoyer/thin
+[event-machine-flush]: https://github.com/oggy/event_machine_flush
 [contact]: mailto:george.ogata@gmail.com
 
 ### Controller
 
- * `when_streaming_template` - defines a callback to be called during a `render`
-   call when a template is streamed. This is *before* the body is rendered, or
-   any data is sent to the client.
+Class methods:
+
+ * `stream` - stream responses for these actions. Takes `:only` or `:except`
+   options, like `before_filter`.
+
+ * `when_streaming_template` - registers a callback to be called during `render`
+   when rendering progressively. This is before the body is rendered, or any
+   data is sent to the client.
+
+Instance methods:
 
-## Example
+ * `render` has been modified to accept a `:stream` option. If true, the
+   response will be streamed, otherwise it won't. This overrides the setting set
+   by the `stream` method above.
+
+### Error Recovery
+
+As mentioned above, headers are sent to the client before view rendering starts,
+which means it's not possible to send an error response in the event of an
+uncaught exception. Instead, the innermost template which raised the error
+simply renders nothing. This has the added advantage of minimizing the impact on
+your visitors, as the rest of the page will render fine.
+
+When an error is swallowed like this, it is passed to an error hander callback,
+which you can set as follows.
+
+    TemplateStreaming.on_streaming_error do |controller, exception|
+      ...
+    end
+
+This is where you should hook in your error notification system. Errors are also
+logged to the application log.
+
+In addition, in development mode, error information is injected into the foot of
+the page. This is presented over the top of the rendered page, so the result
+looks much like when not streaming.
+
+## Streaming Templates Effectively
 
 Conventional wisdom says to put your external stylesheets in the HEAD of your
 page, and your external javascripts at the bottom of the BODY (markup in
 [HAML][haml]):
 
-### `app/views/prelayouts/application.html.haml`
+[haml]: http://haml-lang.com
+
+### `app/views/layouts/application.html.haml`
 
     !!! 5
     %html
@@ -91,65 +150,101 @@ page, and your external javascripts at the bottom of the BODY (markup in
         = stylesheet_link_tag 'one'
         = stylesheet_link_tag 'two'
      - flush
-     = yield
+     %body
+       = yield
+       = javascript_include_tag 'one'
+       = javascript_include_tag 'two'
+
+When streaming, however, you can do better: put the javascripts at the top of
+the page too, and fetch them *asynchronously*. This can be done by appending a
+script tag to the HEAD of the page in a small piece of inline javascript:
 
 ### `app/views/layouts/application.html.haml`
 
+    !!! 5
+    %html
+      %head
+        = stylesheet_link_tag 'one'
+        = stylesheet_link_tag 'two'
+        = javascript_tag do
+          = File.read(Rails.public_path + '/javascripts/get_script.js')
+          $.getScript('#{javascript_path('jquery')}');
+          $.getScript('#{javascript_path('application')}');
     %body
-      = yield
-      = javascript_include_tag 'one'
-      = javascript_include_tag 'two'
+        - flush
+        = yield
+
+### `public/javascripts/get_script.js`
+
+    //
+    // Credit: Sam Cole [https://gist.github.com/364746]
+    //
+    window.$ = {
+      getScript: function(script_src, callback) {
+        var done = false;
+        var head = document.getElementsByTagName("head")[0] || document.documentElement;
+        var script = document.createElement("script");
+        script.src = script_src;
+        script.onload = script.onreadystatechange = function() {
+          if ( !done && (!this.readyState ||
+              this.readyState === "loaded" || this.readyState === "complete") ) {
+            if(callback) callback();
 
-With progressive rendering, however, this could be improved. As [Stoyan Stefanov
-writes][stefanov], you can put your javascripts in the HEAD of your page if you
-fetch them via AJAX and append them to the HEAD of your page dynamically. This
-also reduces the time for the page to become interactive (e.g., scrollable),
-giving an even greater perceived performance boost.
+            // Handle memory leak in IE
+            script.onload = script.onreadystatechange = null;
+            if ( head && script.parentNode ) {
+              head.removeChild( script );
+            }
 
-Of course, rather than using an external library for the AJAX call, we can save
-ourselves a roundtrip by defining a `getScript` function ourselves in a small
-piece of inline javascript. This is done by `define_get_script`
-below. `get_script` then includes a call to this function which fetches the
-script asynchronously, and then appends the script tag to the HEAD.
+            done = true;
+          }
+        };
+        head.insertBefore( script, head.firstChild );
+      }
+    };
+
+If you have inline javascript that depends on the fetched scripts, you'll need
+to delay its execution until the scripts have been run. You can do this by
+wrapping the javascript in a function, with a guard which will delay execution
+until the script is loaded, unless the script has already been loaded. Example:
 
-### `app/views/prelayouts/application.html.haml`
+### Layout
 
     !!! 5
     %html
       %head
-        = define_get_script
         = stylesheet_link_tag 'one'
         = stylesheet_link_tag 'two'
-        = get_script 'one'
-        = get_script 'two'
-     - flush
-     = yield
-
-### `app/views/layouts/application.html.haml`
-
+        = javascript_tag do
+          = File.read(Rails.public_path + '/javascripts/get_script.js')
+          $.getScript('#{javascript_path('jquery')}', function() {
+            window.script_loaded = 1;
+
+            // If the inline code has been loaded (but not yet run), run it
+            // now. Otherwise, it will be run immediately when it's available.
+            if (window.inline)
+              inline();
+          });
     %body
-      = yield
+        - flush
+        = yield
 
-### `app/helpers/application_helper.rb`
+### View
 
-    module ApplicationHelper
-      def define_get_script
-        javascript_tag do
-          File.read(Rails.public_path + '/javascripts/get_script.js')
-        end
-      end
+    - javascript_tag do
+      window.inline() {
+        // ... inline javascript code ...
+      }
 
-      def get_script(url)
-        javascript_tag do
-          "$.getScript('#{javascript_path(url)}');"
-        end
-      end
-    end
+      // If the script is already loaded, run it now. Otherwise, the callback
+      // above will run it after the script is loaded.
+      if (window.script_loaded)
+        inline();
 
-### `public/javascripts/get_script.js`
+### In `public/javascripts/get_script.js`
 
     //
-    // Written by Sam Cole. See http://gist.github.com/364746 for more info.
+    // Credit: Sam Cole [https://gist.github.com/364746]
     //
     window.$ = {
       getScript: function(script_src, callback) {
@@ -175,20 +270,14 @@ script asynchronously, and then appends the script tag to the HEAD.
       }
     };
 
-The second profile was created using this code.
-
-[haml]: http://haml-lang.com
-[stefanov]: http://www.yuiblog.com/blog/2008/07/22/non-blocking-scripts
-[get-script]: http://gist.github.com/364746
-
-## Note on Patches/Pull Requests
+## Contributing
 
- * Bug reports: http://github.com/oggy/template_streaming/issues
- * Source: http://github.com/oggy/template_streaming
+ * [Bug reports](https://github.com/oggy/template_streaming/issues)
+ * [Source](https://github.com/oggy/template_streaming)
  * Patches: Fork on Github, send pull request.
-   * Ensure patch includes tests.
+   * Include tests where practical.
    * Leave the version alone, or bump it in a separate commit.
 
 ## Copyright
 
-Copyright (c) 2010 George Ogata. See LICENSE for details.
+Copyright (c) George Ogata. See LICENSE for details.

data/Rakefile

@@ -1,22 +1 @@
-gem 'ritual'
 require 'ritual'
-
-spec_task :spec do |t|
-  t.libs << 'lib' << 'spec'
-  t.spec_files = FileList['spec/**/*_spec.rb']
-end
-
-spec_task :rcov do |t|
-  t.libs << 'lib' << 'spec'
-  t.pattern = 'spec/**/*_spec.rb'
-  t.rcov = true
-end
-
-rdoc_task do |t|
-  t.rdoc_dir = 'rdoc'
-  t.title = "Template Streaming #{version}"
-  t.rdoc_files.include('README*')
-  t.rdoc_files.include('lib/**/*.rb')
-end
-
-task :default => :spec

data/lib/template_streaming.rb

@@ -5,43 +5,17 @@ module TemplateStreaming
         send "#{key}=", value
       end
     end
-
-    #
-    # If true, always reference the flash before returning from the
-    # action when rendering progressively.
-    #
-    # This is required for the flash to work with progressive
-    # rendering, but unlike standard Rails behavior, will cause the
-    # flash to be swept even if it's never referenced in the
-    # views. This usually isn't an issue, as flash messages are
-    # typically rendered in the layout, causing a reference anyway.
-    #
-    # Default: true.
-    #
-    attr_accessor :autosweep_flash
-
-    #
-    # If true, always set the authenticity token before returning from
-    # the action when rendering progressively.
-    #
-    # This is required for the authenticity token to work with
-    # progressive rendering, but unlike standard Rails behavior, will
-    # cause the token to be set (and thus the session updated) even if
-    # it's never referenced in views.
-    #
-    # Default: true.
-    #
-    attr_accessor :set_authenticity_token
   end
 
-  self.autosweep_flash = true
-  self.set_authenticity_token = true
+  STREAMING_KEY = 'template_streaming.streaming'.freeze
 
   module Controller
     def self.included(base)
       base.class_eval do
+        extend ClassMethods
         alias_method_chain :render, :template_streaming
         alias_method_chain :render_to_string, :template_streaming
+        alias_method_chain :flash, :template_streaming
         helper_method :flush, :push
 
         include ActiveSupport::Callbacks
@@ -49,11 +23,39 @@ module TemplateStreaming
       end
     end
 
+    module ClassMethods
+      def stream(options={})
+        before_filter :action_streams, options
+      end
+    end
+
+    def action_streams
+      @action_streams = true
+    end
+
+    def action_streams?
+      @action_streams
+    end
+
     def render_with_template_streaming(*args, &block)
-      with_template_streaming_condition(*args) do |condition|
-        if condition
+      options = args.first { |a| a.is_a?(Hash) }
+      if options && options.size == 1 && options.key?(:stream)
+        # Need to set the default values, since the standard #render won't.
+        options[:template] = default_template
+        options[:layout] = true
+      end
+      push_render_stack_frame do |stack_height|
+        if start_streaming_template?(stack_height, *args)
+          @streaming_template = true
+          @template.streaming_template = true
           @performed_render = true
-          @streaming_body = StreamingBody.new(progressive_rendering_threshold) do
+          @streaming_body = StreamingBody.new(template_streaming_threshold) do
+            cookies.freeze
+            if self.class.session_store.sent_with_headers?
+              session.freeze
+              flash.freeze
+            end
+
             @performed_render = false
             last_piece = render_without_template_streaming(*args, &block)
             # The original render will clobber our response.body, so
@@ -62,33 +64,44 @@ module TemplateStreaming
           end
           response.body = @streaming_body
           response.prepare!
-          flash if TemplateStreaming.autosweep_flash
-          form_authenticity_token if TemplateStreaming.set_authenticity_token
-          run_callbacks :when_streaming_template
+          form_authenticity_token  # generate now
 
-          # Normally, @_flash is removed after #perform_action, which
-          # means calling #flash in the view would cause a new
-          # FlashHash to be constructed. On top of that, the flash is
-          # swept on construction, which results in sweeping the flash
-          # twice, obliterating its contents.
+          # Normally, the flash is swept on first reference. This
+          # means we need to ensure it's referenced before the session
+          # is persisted. In the case of the cookie store, that's when
+          # the headers are sent, so we force a reference now.
           #
-          # So, we preserve the flash here under a different ivar, and
-          # override the #flash helper to return it.
-          if defined?(@_flash)
-            @template_streaming_flash = @_flash
-          end
+          # But alas, that's not all. @_flash is removed after
+          # #perform_action, which means calling #flash in the view
+          # would cause the flash to be referenced again, sweeping the
+          # flash a second time. To prevent this, we preserve the
+          # flash in a separate ivar, and patch #flash to return this
+          # if we're streaming.
+          #
+          flash  # ensure sweep
+          @template_streaming_flash = @_flash
+          request.env[STREAMING_KEY] = true
+
+          run_callbacks :when_streaming_template
         else
           render_without_template_streaming(*args, &block)
         end
       end
     end
 
+    # Mark the case when it's a layout for a toplevel render. This is
+    # done here, as it's called after the option wrangling in
+    # AC::Base#render, and nowhere else.
+    def pick_layout(options)
+      result = super
+      options[:toplevel_render_with_layout] = true if result
+      result
+    end
+
     # Override to ensure calling render_to_string from a helper
     # doesn't trigger template streaming.
     def render_to_string_with_template_streaming(*args, &block) # :nodoc
-      # Ensure renders within a render_to_string aren't considered
-      # top-level.
-      with_template_streaming_condition do
+      push_render_stack_frame do
         render_to_string_without_template_streaming(*args, &block)
       end
     end
@@ -98,7 +111,7 @@ module TemplateStreaming
     # immediately.
     #
     def flush
-      unless @template.output_buffer.nil?
+      if @streaming_body && !@template.output_buffer.nil?
         push @template.output_buffer.slice!(0..-1)
       end
     end
@@ -107,35 +120,50 @@ module TemplateStreaming
     # Push the given data to the client immediately.
     #
     def push(data)
-      @streaming_body.push(data)
-      flush_thin
+      if @streaming_body
+        @streaming_body.push(data)
+        flush_thin
+      end
     end
 
     def template_streaming_flash # :nodoc:
       @template_streaming_flash
     end
 
+    def streaming_template?
+      @streaming_template
+    end
+
     private # --------------------------------------------------------
 
-    #
-    # Yield true if we should intercept this render call, false
-    # otherwise.
-    #
-    def with_template_streaming_condition(*args)
+    def push_render_stack_frame
       @render_stack_height ||= 0
       @render_stack_height += 1
       begin
-        # Only install our StreamingBody in the toplevel #render call.
-        @render_stack_height == 1 or
-          return yield(false)
+        yield @render_stack_height
+      ensure
+        @render_stack_height -= 1
+      end
+    end
+
+    def start_streaming_template?(render_stack_height, *render_args)
+      render_stack_height == 1 or
+        return false
+
+      return false if rendered_action_cache
 
-        if (options = args.last).is_a?(Hash)
-          yield((UNSTREAMABLE_KEYS & options.keys).empty?)
+      (render_options = render_args.last).is_a?(Hash) or
+        render_options = {}
+
+      if !(UNSTREAMABLE_KEYS & render_options.keys).empty? || render_args.first == :update
+        false
+      else
+        explicit_option = render_options[:stream]
+        if explicit_option.nil?
+          action_streams?
         else
-          yield(args.first != :update)
+          explicit_option
         end
-      ensure
-        @render_stack_height -= 1
       end
     end
 
@@ -145,7 +173,7 @@ module TemplateStreaming
     # The number of bytes that must be received by the client before
     # anything will be rendered.
     #
-    def progressive_rendering_threshold
+    def template_streaming_threshold
       content_type = response.header['Content-type']
       content_type.nil? || content_type =~ %r'\Atext/html' or
         return 0
@@ -170,6 +198,15 @@ module TemplateStreaming
       connection = request.env['template_streaming.thin_connection'] and
         EventMachineFlush.flush(connection)
     end
+
+    def flash_with_template_streaming # :nodoc:
+      if defined?(@template_streaming_flash)
+        # Flash has been swept - don't use the standard #flash or it'll sweep again.
+        @template_streaming_flash
+      else
+        flash_without_template_streaming
+      end
+    end
   end
 
   # Only prepare once.
@@ -196,53 +233,83 @@ module TemplateStreaming
 
   module View
     def self.included(base)
+      base.alias_method_chain :render, :template_streaming
       base.alias_method_chain :_render_with_layout, :template_streaming
-      base.alias_method_chain :flash, :template_streaming
     end
 
-    def _render_with_layout_with_template_streaming(options, local_assigns, &block)
-      with_prelayout prelayout_for(options), local_assigns do
-        _render_with_layout_without_template_streaming(options, local_assigns, &block)
+    def render_with_template_streaming(*args, &block)
+      options = args.first
+      if streaming_template? && options.is_a?(Hash)
+        # These branches exist to handle the case where AC::Base#render calls
+        # AV::Base#render for rendering a partial with a layout. AC::Base
+        # renders the partial then the layout separately, but we need to render
+        # them together, in the reverse order (layout first). We do this by
+        # standard-rendering the layout with a block that renders the partial.
+        if options[:toplevel_render_with_layout] && (partial = options[:partial])
+          # Don't render yet - we need to do the layout first.
+          options.delete(:toplevel_render_with_layout)
+          return DeferredPartialRender.new(args)
+        elsif options[:text].is_a?(DeferredPartialRender)
+          render = options.delete(:text)
+          # We patch the case of rendering :partial with :layout while
+          # streaming in _render_with_layout.
+          return render(render.args.first.merge(:layout => options[:layout]))
+        end
       end
+      render_without_template_streaming(*args, &block)
+    end
+
+    DeferredPartialRender = Struct.new(:args)
+
+    attr_writer :streaming_template
+
+    def streaming_template?
+      @streaming_template
     end
 
-    def with_prelayout(prelayout, locals, &block)
-      if prelayout
+    def _render_with_layout_with_template_streaming(options, local_assigns, &block)
+      if !streaming_template?
+        _render_with_layout_without_template_streaming(options, local_assigns, &block)
+      elsif block_given?
+        # The standard method doesn't properly restore @_proc_for_layout. Do it ourselves.
+        original_proc_for_layout = @_proc_for_layout
         begin
-          @_proc_for_layout = lambda do
-            # nil out @_proc_for_layout else rendering with the layout will call it again.
-            @_proc_for_layout, original_proc_for_layout = nil, @_proc_for_layout
-            begin
-              block.call
-            ensure
-              @_proc_for_layout = original_proc_for_layout
-            end
-          end
-          render(:file => prelayout, :locals => locals)
+          _render_with_layout_without_template_streaming(options, local_assigns, &block)
         ensure
-          @_proc_for_layout = nil
+          @_proc_for_layout = original_proc_for_layout
+        end
+      elsif options[:layout].is_a?(ActionView::Template)
+        # Toplevel render call, from the controller.
+        layout = options.delete(:layout)
+        with_render_proc_for_layout(options) do
+          render(options.merge(:file => layout.path_without_format_and_extension))
         end
       else
-        yield
+        layout = options.delete(:layout)
+        with_render_proc_for_layout(options) do
+          if (options[:inline] || options[:file] || options[:text])
+            render(:file => layout, :locals => local_assigns)
+          else
+            render(options.merge(:partial => layout))
+          end
+        end
       end
     end
 
-    def prelayout_for(options)
-      layout = options[:layout] or
-        return nil
-      # Views can call #render with :layout to render a layout
-      # *partial* which we don't want to interfere with. Only the
-      # interlal toplevel #render calls :layout with an
-      # ActionView::Template
-      layout.is_a?(ActionView::Template) or
-        return nil
-      view_paths.find_template('pre' + layout.path_without_format_and_extension, layout.format)
-    rescue ActionView::MissingTemplate
-    end
-
-    def flash_with_template_streaming # :nodoc:
-      # Override ActionView::Base#flash to prevent a double-sweep.
-      controller.instance_eval { @template_streaming_flash || flash }
+    def with_render_proc_for_layout(options)
+      original_proc_for_layout = @_proc_for_layout
+      @_proc_for_layout = lambda do |*args|
+        if args.empty?
+          render(options)
+        else
+          instance_variable_get(:"@content_for_#{args.first}")
+        end
+      end
+      begin
+        yield
+      ensure
+        @_proc_for_layout = original_proc_for_layout
+      end
     end
   end
 
@@ -275,10 +342,24 @@ module TemplateStreaming
     end
   end
 
+  module AbstractSessionStoreExtension
+    def sent_with_headers?
+      false
+    end
+  end
+
+  module CookieSessionStoreExtension
+    def sent_with_headers?
+      true
+    end
+  end
+
   ActionView::Base.send :include, View
   ActionController::Base.send :include, Controller
   ActionController::Response.send :include, Response
   ActionController::Dispatcher.middleware.insert 0, Rack::Chunked
+  ActionController::Session::AbstractStore.extend AbstractSessionStoreExtension
+  ActionController::Session::CookieStore.extend CookieSessionStoreExtension
 end
 
 # Please let there be a better way to do this...
@@ -311,3 +392,7 @@ if defined?(Thin)
     end
   end
 end
+
+require 'template_streaming/error_recovery'
+require 'template_streaming/caching'
+require 'template_streaming/autoflushing'