Rails concern for streaming response

streaming_response.rb

# frozen_string_literal: true

require "timeout"

module StreamingResponse
  extend ActiveSupport::Concern

  included do
    include ActionController::Live
  end

  private

  # Generic streaming wrapper with error handling
  # @param timeout_seconds [Integer] Maximum time to wait for the stream (default: 300)
  def stream_response(timeout_seconds: 300)
    setup_streaming_headers

    begin
      Timeout.timeout(timeout_seconds) do
        yield if block_given?
      end
    rescue Timeout::Error
      Rails.logger.error "Stream timed out after #{timeout_seconds} seconds"
    rescue IOError, ActionController::Live::ClientDisconnected => e
      Rails.logger.info "Client disconnected: #{e.message}"
    rescue => e
      Rails.logger.error "Stream error: #{e.message}"
    ensure
      cleanup_stream
    end
  end

  # Write raw data to stream - completely format agnostic
  # @param data [String] Any data to stream
  def stream_write(data)
    response.stream.write(data)
  end

  def setup_streaming_headers
    response.headers["Content-Type"] = "text/plain; charset=utf-8"
    response.headers["Cache-Control"] = "no-cache"
    response.headers["Connection"] = "keep-alive"
    response.headers["X-Accel-Buffering"] = "no" # Disable nginx buffering
  end

  def cleanup_stream
    response.stream.close if response.stream && !response.stream.closed?
  rescue IOError
    # Stream already closed
  end
end