style_capsule 1.4.0 → 2.0.0

data/lib/style_capsule/css_file_writer.rb

@@ -75,6 +75,8 @@ module StyleCapsule
     DEFAULT_OUTPUT_DIR = "app/assets/builds/capsules"
     # Fallback directory for when default location is read-only (absolute path)
     FALLBACK_OUTPUT_DIR = "/tmp/style_capsule"
+    # Positive cache for resolved asset-relative paths (avoids repeated File.exist? on hot path)
+    FILE_PATH_CACHE_MAX = 512
 
     class << self
       attr_accessor :output_dir, :filename_pattern, :enabled, :fallback_dir
@@ -99,23 +101,24 @@ module StyleCapsule
       #     filename_pattern: ->(klass, capsule) { "#{klass.name.underscore}-#{capsule}.css" }
       #   )
       def configure(output_dir: nil, filename_pattern: nil, enabled: true, fallback_dir: nil)
+        clear_file_path_hit_cache!
         @enabled = enabled
+        @output_dir = resolve_output_dir(output_dir)
+        @fallback_dir = resolve_path(fallback_dir, FALLBACK_OUTPUT_DIR)
+        @filename_pattern = filename_pattern || default_filename_pattern
+      end
 
-        @output_dir = if output_dir
-          output_dir.is_a?(Pathname) ? output_dir : Pathname.new(output_dir.to_s)
-        elsif rails_available?
-          Rails.root.join(DEFAULT_OUTPUT_DIR)
-        else
-          Pathname.new(DEFAULT_OUTPUT_DIR)
-        end
+      def resolve_output_dir(output_dir)
+        return resolve_path(output_dir, DEFAULT_OUTPUT_DIR) if output_dir
+        return Rails.root.join(DEFAULT_OUTPUT_DIR) if rails_available?
 
-        @fallback_dir = if fallback_dir
-          fallback_dir.is_a?(Pathname) ? fallback_dir : Pathname.new(fallback_dir.to_s)
-        else
-          Pathname.new(FALLBACK_OUTPUT_DIR)
-        end
+        Pathname.new(DEFAULT_OUTPUT_DIR)
+      end
 
-        @filename_pattern = filename_pattern || default_filename_pattern
+      def resolve_path(dir, default_path)
+        return Pathname.new(default_path) unless dir
+
+        dir.is_a?(Pathname) ? dir : Pathname.new(dir.to_s)
       end
 
       # Write CSS content to file
@@ -124,6 +127,7 @@ module StyleCapsule
       # @param component_class [Class] Component class that generated the CSS
       # @param capsule_id [String] Capsule ID for the component
       # @return [String, nil] Relative file path (for stylesheet_link_tag) or nil if disabled/failed
+      # rubocop:disable Metrics/AbcSize -- writes asset file with fallback and instrumentation paths
       def write_css(css_content:, component_class:, capsule_id:)
         return nil unless enabled?
 
@@ -194,13 +198,11 @@ module StyleCapsule
         # Return relative path for stylesheet_link_tag
         # Path should be relative to app/assets
         # Handle case where output directory is not under rails_assets_root (e.g., in tests)
-        begin
-          file_path.relative_path_from(rails_assets_root).to_s.gsub(/\.css$/, "")
-        rescue ArgumentError
-          # If paths don't share a common prefix (e.g., in tests), return just the filename
-          filename.gsub(/\.css$/, "")
-        end
+        rel = asset_logical_path(file_path, filename)
+        remember_file_path_hit("#{component_class.name}:#{capsule_id}", rel)
+        rel
       end
+      # rubocop:enable Metrics/AbcSize
 
       # Check if file exists for given component and capsule
       #
@@ -223,6 +225,10 @@ module StyleCapsule
       def file_path_for(component_class:, capsule_id:)
         return nil unless enabled?
 
+        cache_key = "#{component_class.name}:#{capsule_id}"
+        hit = file_path_hit_cache_mutex.synchronize { file_path_hit_cache[cache_key] }
+        return hit if hit
+
         filename = generate_filename(component_class, capsule_id)
         file_path = output_directory.join(filename)
 
@@ -230,12 +236,9 @@ module StyleCapsule
 
         # Return relative path for stylesheet_link_tag
         # Handle case where output directory is not under rails_assets_root (e.g., in tests)
-        begin
-          file_path.relative_path_from(rails_assets_root).to_s.gsub(/\.css$/, "")
-        rescue ArgumentError
-          # If paths don't share a common prefix (e.g., in tests), return just the filename
-          filename.gsub(/\.css$/, "")
-        end
+        rel = asset_logical_path(file_path, filename)
+        remember_file_path_hit(cache_key, rel)
+        rel
       end
 
       # Ensure output directory exists
@@ -264,6 +267,7 @@ module StyleCapsule
       def clear_files
         return unless enabled?
 
+        clear_file_path_hit_cache!
         dir = output_directory
         return unless Dir.exist?(dir)
 
@@ -285,6 +289,39 @@ module StyleCapsule
 
       private
 
+      def file_path_hit_cache
+        @file_path_hit_cache ||= {}
+      end
+
+      def file_path_hit_cache_mutex
+        @file_path_hit_cache_mutex ||= Mutex.new
+      end
+
+      def remember_file_path_hit(cache_key, relative_path)
+        file_path_hit_cache_mutex.synchronize do
+          file_path_hit_cache[cache_key] = relative_path
+          while file_path_hit_cache.size > FILE_PATH_CACHE_MAX
+            file_path_hit_cache.shift
+          end
+        end
+      end
+
+      def clear_file_path_hit_cache!
+        file_path_hit_cache_mutex.synchronize do
+          file_path_hit_cache.clear
+        end
+      end
+
+      # Logical path for stylesheet_link_tag (no parent segments; safe for AssetPath validation)
+      def asset_logical_path(file_path, filename)
+        candidate = file_path.relative_path_from(rails_assets_root).to_s.gsub(/\.css$/, "")
+        return candidate unless candidate.split("/").any? { |segment| segment == ".." }
+
+        filename.gsub(/\.css$/, "")
+      rescue ArgumentError
+        filename.gsub(/\.css$/, "")
+      end
+
       # Get output directory (with default)
       def output_directory
         @output_dir ||= if rails_available?

data/lib/style_capsule/css_processor.rb

@@ -40,6 +40,7 @@ module StyleCapsule
     # @param component_class [Class, String, nil] Optional component class for instrumentation
     # @return [String] CSS with scoped selectors
     # @raise [ArgumentError] If CSS content exceeds maximum size or capsule_id is invalid
+    # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity -- selector patching walks CSS rules
     def self.scope_selectors(css_string, capsule_id, component_class: nil)
       return css_string if css_string.nil? || css_string.strip.empty?
 
@@ -61,8 +62,8 @@ module StyleCapsule
         css = css_string.dup
         capsule_attr = %([data-capsule="#{capsule_id}"])
 
-        # Strip CSS comments to avoid interference with selector matching
-        # Simple approach: remove /* ... */ comments (including multi-line)
+        # Strip CSS comments so they do not interfere with selector matching.
+        # Comments are removed from the output (typical for production CSS).
         css_without_comments = strip_comments(css)
 
         # Process CSS rule by rule
@@ -108,13 +109,10 @@ module StyleCapsule
           "#{prefix}#{whitespace}#{scoped_selectors}#{opening_brace}"
         end
 
-        # Restore comments in their original positions
-        # Since we stripped comments, we need to put them back
-        # For simplicity, we'll just return the processed CSS without comments
-        # (comments are typically removed in production CSS anyway)
         css_without_comments
       end
     end
+    # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
 
     # Scope CSS using CSS nesting (wraps entire CSS in [data-capsule] { ... })
     #

data/lib/style_capsule/head_injection_middleware.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module StyleCapsule
+  # Injects stylesheets registered during body rendering into +<head>+ on the same request.
+  #
+  # Layouts call +stylesheet_registry_tags+ before the body, so +register_stylesheet+ in
+  # components runs too late for that call. This middleware appends pending link/style tags
+  # immediately before +</head>+ once the full HTML response is available.
+  class HeadInjectionMiddleware
+    HTML_CONTENT_TYPE = %r{\Atext/html}i
+    SUCCESS_STATUS = (200..299)
+
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      status, headers, response = @app.call(env)
+      return [status, headers, response] unless inject_response?(status, headers, response)
+      return [status, headers, response] unless StylesheetRegistry.pending_head_stylesheets?
+
+      maybe_inject_head_styles(status, headers, response, env)
+    end
+
+    private
+
+    def maybe_inject_head_styles(status, headers, response, env)
+      body = read_response_body(response)
+      return [status, headers, response] if body.empty?
+
+      view_context = view_context_for(env)
+      injected_body = StylesheetRegistry.inject_pending_head_stylesheets(body, view_context)
+      return [status, headers, response] if injected_body.equal?(body)
+
+      headers = headers.dup
+      headers["Content-Length"] = injected_body.bytesize.to_s if headers.key?("Content-Length")
+
+      [status, headers, [injected_body]]
+    end
+
+    def inject_response?(status, headers, response)
+      return false unless SUCCESS_STATUS.cover?(status)
+      return false unless html_content_type?(headers["Content-Type"])
+      return false if chunked_response?(headers)
+      return false if response.respond_to?(:to_ary) && response.to_ary == [""]
+
+      true
+    end
+
+    def chunked_response?(headers)
+      transfer_encoding = headers["Transfer-Encoding"]
+      transfer_encoding&.match?(/chunked/i)
+    end
+
+    def html_content_type?(content_type)
+      return false if content_type.blank?
+
+      content_type.split(";", 2).first.to_s.strip.match?(HTML_CONTENT_TYPE)
+    end
+
+    def read_response_body(response)
+      parts = []
+      response.each { |part| parts << part.to_s }
+      parts.join
+    end
+
+    def view_context_for(env)
+      controller = env["action_controller.instance"]
+      controller&.view_context
+    end
+  end
+end

data/lib/style_capsule/helper.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "digest/sha1"
+require_relative "helper_scope_cache"
 # ActiveSupport string extensions are conditionally required in lib/style_capsule.rb
 
 module StyleCapsule
@@ -24,6 +25,8 @@ module StyleCapsule
   # The CSS will be automatically scoped and content wrapped in a scoped div.
   # The <style> tag will be extracted and processed separately.
   module Helper
+    include HelperScopeCache
+
     # Maximum HTML content size (10MB) to prevent DoS attacks
     MAX_HTML_SIZE = 10_000_000
     # Generate capsule ID based on caller location for uniqueness
@@ -36,16 +39,7 @@ module StyleCapsule
 
     # Scope CSS content and return scoped CSS
     def scope_css(css_content, capsule_id)
-      # Use thread-local cache to avoid reprocessing
-      cache_key = "style_capsule_#{capsule_id}"
-
-      if Thread.current[cache_key]
-        return Thread.current[cache_key]
-      end
-
-      scoped_css = CssProcessor.scope_selectors(css_content, capsule_id)
-      Thread.current[cache_key] = scoped_css
-      scoped_css
+      scope_css_with_bounded_cache(css_content, capsule_id)
     end
 
     # ERB helper: automatically wraps content in scoped div and processes CSS
@@ -83,7 +77,14 @@ module StyleCapsule
     #   <%= style_capsule(css_content, capsule_id: "test-123") do %>
     #     <div class="section">Content</div>
     #   <% end %>
-    def style_capsule(css_content = nil, capsule_id: nil, &content_block)
+    #
+    # Or with tag option:
+    #   <%= style_capsule(tag: :section) do %>
+    #     <style>.section { color: red; }</style>
+    #     <div class="section">Content</div>
+    #   <% end %>
+    # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity -- ERB helper extracts styles and wraps scoped markup
+    def style_capsule(css_content = nil, capsule_id: nil, tag: :div, &content_block)
       html_content = nil
 
       # If CSS content is provided as argument, use it
@@ -96,17 +97,9 @@ module StyleCapsule
           raise ArgumentError, "HTML content exceeds maximum size of #{MAX_HTML_SIZE} bytes (got #{full_content.bytesize} bytes)"
         end
 
-        # Extract <style> tags from content
-        # Note: Pattern uses non-greedy matching (.*?) to minimize backtracking
-        # Size limit (MAX_HTML_SIZE) mitigates ReDoS risk from malicious input
-        style_match = full_content.match(/<style[^>]*>(.*?)<\/style>/m)
-        if style_match
-          css_content = style_match[1]
-          # Use sub instead of gsub to only remove first occurrence (reduces backtracking)
-          html_content = full_content.sub(/<style[^>]*>.*?<\/style>/m, "").strip
-        else
-          # No style tag found, treat entire content as HTML
-          css_content = nil
+        # Extract every <style> block; combine CSS, strip tags from markup
+        html_content, css_content = extract_styles_from_markup(full_content)
+        if css_content.nil?
           html_content = full_content
         end
       elsif css_content && block_given?
@@ -129,10 +122,11 @@ module StyleCapsule
 
       # Render style tag and wrapped content
       style_tag = content_tag(:style, raw(scoped_css), type: "text/css")
-      wrapped_content = content_tag(:div, raw(html_content), data: {capsule: capsule_id})
+      wrapped_content = content_tag(tag, raw(html_content), data: {capsule: capsule_id})
 
       (style_tag + wrapped_content).html_safe
     end
+    # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
 
     # Register a stylesheet file for head rendering
     #

data/lib/style_capsule/helper_scope_cache.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+require "strscan"
+
+module StyleCapsule
+  # Bounded per-thread cache for scoped CSS from ERB / standalone helpers (avoids unbounded Thread.current growth).
+  module HelperScopeCache
+    MAX_SCOPE_CACHE_ENTRIES = 256
+    STYLE_OPEN = %r{<style[^>]*>}im
+    STYLE_CLOSE = %r{</style>}im
+
+    private
+
+    # @param full_html [String]
+    # @return [Array(String, String, nil)] html_without_styles, combined_css (nil if none)
+    def extract_styles_from_markup(full_html)
+      scanner = StringScanner.new(full_html)
+      html_out = +""
+      segments = []
+
+      until scanner.eos?
+        break unless collect_style_segment!(scanner, html_out, segments)
+      end
+
+      combined = segments.empty? ? nil : segments.join("\n\n")
+      [html_out.strip, combined]
+    end
+
+    def collect_style_segment!(scanner, html_out, segments)
+      text = scanner.scan_until(STYLE_OPEN)
+      if text.nil?
+        html_out << scanner.rest
+        return false
+      end
+
+      html_out << text.sub(STYLE_OPEN, "")
+
+      css = scanner.scan_until(STYLE_CLOSE)
+      if css.nil?
+        html_out << scanner.rest
+        return false
+      end
+
+      segments << css.sub(STYLE_CLOSE, "")
+      true
+    end
+
+    def scope_css_with_bounded_cache(css_content, capsule_id)
+      fingerprint = Digest::SHA1.hexdigest(css_content.to_s)
+      cache_key = "style_capsule_#{capsule_id}_#{fingerprint}"
+
+      bucket = Thread.current[:style_capsule_scope_cache] ||= {order: [], hash: {}}
+      hash = bucket[:hash]
+      order = bucket[:order]
+
+      return hash[cache_key] if hash.key?(cache_key)
+
+      scoped_css = CssProcessor.scope_selectors(css_content, capsule_id)
+      evict_scope_cache_if_full!(hash, order)
+      hash[cache_key] = scoped_css
+      order << cache_key
+      scoped_css
+    end
+
+    def evict_scope_cache_if_full!(hash, order)
+      while order.size >= MAX_SCOPE_CACHE_ENTRIES
+        old = order.shift
+        hash.delete(old)
+      end
+    end
+  end
+end

data/lib/style_capsule/instrumentation.rb

@@ -12,6 +12,7 @@ module StyleCapsule
   # @example Subscribing to CSS processing events
   #   ActiveSupport::Notifications.subscribe("style_capsule.css_processor.scope") do |name, start, finish, id, payload|
   #     duration_ms = (finish - start) * 1000
+  #     # input_size is derived from :css_content when subscribers exist
   #     input_size = payload[:input_size]
   #     Rails.logger.info "CSS scoped in #{duration_ms}ms, input: #{input_size} bytes"
   #   end
@@ -31,7 +32,7 @@ module StyleCapsule
     #
     # @return [Boolean]
     def self.available?
-      defined?(ActiveSupport::Notifications)
+      !!defined?(ActiveSupport::Notifications)
     end
 
     # Instrument an operation with automatic timing and size metrics
@@ -52,10 +53,11 @@ module StyleCapsule
     #     component_class: component_class.name,
     #     capsule_id: capsule_id,
     #     strategy: :selector_patching,
-    #     input_size: css_content.bytesize
+    #     css_content: css_content
     #   ) do
     #     CssProcessor.scope_selectors(css_content, capsule_id)
     #   end
+    #   # When subscribers exist, :input_size is set from :css_content before the block runs.
     def self.instrument(event_name, payload = {}, &block)
       return yield unless available?
 
@@ -65,14 +67,7 @@ module StyleCapsule
         return yield
       end
 
-      # Calculate input size if not provided but css_content/input is in payload
-      unless payload[:input_size]
-        if payload[:css_content]
-          payload[:input_size] = payload[:css_content].bytesize
-        elsif payload[:input]
-          payload[:input_size] = payload[:input].bytesize
-        end
-      end
+      apply_input_size!(payload)
 
       # Use ActiveSupport::Notifications.instrument which automatically:
       # - Measures duration using monotonic time (available in event.duration)
@@ -87,6 +82,14 @@ module StyleCapsule
       result
     end
 
+    def self.apply_input_size!(payload)
+      return if payload[:input_size]
+
+      source = payload[:css_content] || payload[:input]
+      payload[:input_size] = source.bytesize if source
+    end
+    private_class_method :apply_input_size!
+
     # Instrument an event without a block (for events that don't have duration)
     #
     # @param event_name [String] Event name
@@ -115,23 +118,20 @@ module StyleCapsule
     # @param strategy [Symbol] Scoping strategy (:selector_patching or :nesting)
     # @param component_class [Class, String] Component class
     # @param capsule_id [String] Capsule ID
-    # @param css_content [String] CSS content (for size calculation)
+    # @param css_content [String] CSS content (`:input_size` derived lazily when subscribed)
     # @yield The CSS processing operation
     # @return [String] Processed CSS
     def self.instrument_css_processing(strategy:, component_class:, capsule_id:, css_content:, &block)
       component_name = component_class.is_a?(Class) ? component_class.name : component_class.to_s
-      input_size = css_content.bytesize
 
       instrument(
         "style_capsule.css_processor.scope",
         strategy: strategy,
         component_class: component_name,
         capsule_id: capsule_id,
-        input_size: input_size
+        css_content: css_content
       ) do
         result = yield
-        # Output size will be calculated by subscribers if needed
-        # They can access it via: result.bytesize
         result
       end
     end

data/lib/style_capsule/phlex_helper.rb

@@ -16,7 +16,7 @@ module StyleCapsule
   # Usage in Phlex components:
   #   class MyComponent < ApplicationComponent
   #     include StyleCapsule::PhlexHelper
-  #     styles_namespace :user  # Set default namespace
+  #     style_capsule namespace: :user  # Set default namespace for register_stylesheet
   #
   #     def view_template
   #       register_stylesheet("stylesheets/user/my_component")  # Uses :user namespace automatically
@@ -33,7 +33,7 @@ module StyleCapsule
     #     div { "Content" }
     #   end
     #
-    # If the component has a default namespace set via styles_namespace or stylesheet_registry,
+    # If the component has a default namespace set via style_capsule or stylesheet_registry,
     # it will be used automatically when namespace is not explicitly provided.
     #
     # @param file_path [String] Path to stylesheet (relative to app/assets/stylesheets)
@@ -60,18 +60,13 @@ module StyleCapsule
     # @return [void] Renders stylesheet tags via raw
     def stylesheet_registry_tags(namespace: nil)
       output = StyleCapsule::StylesheetRegistry.render_head_stylesheets(view_context, namespace: namespace)
-      # Phlex's raw() requires the object to be marked as safe
-      # Use Phlex's safe() if available, otherwise fall back to html_safe for test doubles
-      # The output from render_head_stylesheets is already html_safe (SafeBuffer)
       output_string = output.to_s
 
       if respond_to?(:safe)
-        # Real Phlex component - use raw() for rendering
-        safe_content = safe(output_string)
-        raw(safe_content)
+        raw(safe(output_string))
       end
 
-      # Always return the output string for testing/compatibility
+      # Phlex `raw` may return nil when writing to the buffer; callers/tests expect the HTML string.
       output_string
     end
   end

data/lib/style_capsule/railtie.rb

@@ -4,6 +4,10 @@ module StyleCapsule
   # Railtie to automatically include StyleCapsule helpers in Rails
   if defined?(Rails::Railtie)
     class Railtie < Rails::Railtie
+      config.style_capsule = ActiveSupport::OrderedOptions.new
+      config.style_capsule.run_on_precompile = true
+      config.style_capsule.head_injection_middleware = true
+
       # Automatically include ERB helper in ActionView::Base (standard Rails pattern)
       # This makes helpers available in all ERB templates automatically
       ActiveSupport.on_load(:action_view) do
@@ -53,6 +57,13 @@ module StyleCapsule
         load File.expand_path("../tasks/style_capsule.rake", __dir__)
       end
 
+      initializer "style_capsule.head_injection_middleware", after: :load_config_initializers do |app|
+        next unless app.config.style_capsule.head_injection_middleware
+
+        require_relative "head_injection_middleware"
+        app.config.middleware.use StyleCapsule::HeadInjectionMiddleware
+      end
+
       # Note: PhlexHelper should be included explicitly in ApplicationComponent
       # or your base Phlex component class, not automatically via Railtie
     end

data/lib/style_capsule/standalone_helper.rb

@@ -2,6 +2,7 @@
 
 require "digest/sha1"
 require "cgi"
+require_relative "helper_scope_cache"
 
 module StyleCapsule
   # Standalone helper module for use without Rails
@@ -27,6 +28,8 @@ module StyleCapsule
   #     "<style>.section { color: red; }</style><div class='section'>Content</div>"
   #   end
   module StandaloneHelper
+    include HelperScopeCache
+
     # Maximum HTML content size (10MB) to prevent DoS attacks
     MAX_HTML_SIZE = 10_000_000
 
@@ -40,16 +43,7 @@ module StyleCapsule
 
     # Scope CSS content and return scoped CSS
     def scope_css(css_content, capsule_id)
-      # Use thread-local cache to avoid reprocessing
-      cache_key = "style_capsule_#{capsule_id}"
-
-      if Thread.current[cache_key]
-        return Thread.current[cache_key]
-      end
-
-      scoped_css = CssProcessor.scope_selectors(css_content, capsule_id)
-      Thread.current[cache_key] = scoped_css
-      scoped_css
+      scope_css_with_bounded_cache(css_content, capsule_id)
     end
 
     # Generate HTML tag without Rails helpers
@@ -59,6 +53,7 @@ module StyleCapsule
     # @param options [Hash] HTML attributes
     # @param block [Proc] Block for tag content
     # @return [String] HTML string
+    # rubocop:disable Metrics/AbcSize -- serializes flat and nested HTML attributes
     def content_tag(tag, content = nil, **options, &block)
       tag_name = tag.to_s
       content = capture(&block) if block_given? && content.nil?
@@ -76,6 +71,7 @@ module StyleCapsule
       attrs = " #{attrs}" unless attrs.empty?
       "<#{tag_name}#{attrs}>#{content}</#{tag_name}>"
     end
+    # rubocop:enable Metrics/AbcSize
 
     # Capture block content (simplified version without Rails)
     #
@@ -114,7 +110,8 @@ module StyleCapsule
     # @param capsule_id [String, nil] Optional capsule ID
     # @param content_block [Proc] Block containing HTML content
     # @return [String] HTML with scoped CSS and wrapped content
-    def style_capsule(css_content = nil, capsule_id: nil, &content_block)
+    # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity -- standalone helper mirrors Rails helper flow
+    def style_capsule(css_content = nil, capsule_id: nil, tag: :div, &content_block)
       html_content = nil
 
       # If CSS content is provided as argument, use it
@@ -127,13 +124,8 @@ module StyleCapsule
           raise ArgumentError, "HTML content exceeds maximum size of #{MAX_HTML_SIZE} bytes (got #{full_content.bytesize} bytes)"
         end
 
-        # Extract <style> tags from content
-        style_match = full_content.match(/<style[^>]*>(.*?)<\/style>/m)
-        if style_match
-          css_content = style_match[1]
-          html_content = full_content.sub(/<style[^>]*>.*?<\/style>/m, "").strip
-        else
-          css_content = nil
+        html_content, css_content = extract_styles_from_markup(full_content)
+        if css_content.nil?
           html_content = full_content
         end
       elsif css_content && block_given?
@@ -156,10 +148,11 @@ module StyleCapsule
 
       # Render style tag and wrapped content
       style_tag = content_tag(:style, raw(scoped_css), type: "text/css")
-      wrapped_content = content_tag(:div, raw(html_content), data: {capsule: capsule_id})
+      wrapped_content = content_tag(tag, raw(html_content), data: {capsule: capsule_id})
 
       html_safe(style_tag + wrapped_content)
     end
+    # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
 
     # Register a stylesheet file for head rendering
     #