search-engine-for-typesense 1.0.0

data/lib/search_engine/instrumentation.rb

@@ -0,0 +1,308 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  # Lightweight wrapper over ActiveSupport::Notifications that standardizes
+  # event names, payload shaping, and timing helpers. Provides a thread-local
+  # context to propagate dispatch metadata (e.g., mode, job_id) from the
+  # dispatcher/job into lower layers such as the indexer.
+  #
+  # Public API is intentionally small:
+  # - {.instrument(event, payload = {}) { |ctx| }}: emit an event (yields mutable ctx)
+  # - {.time(event, base_payload = {}) { |payload| }}: measure duration_ms
+  # - {.with_context(hash) { }}: set per-thread context for nested calls
+  # - {.context}: current shallow context Hash (dup)
+  #
+  # All payloads are duped, nil values are pruned, and keys are symbolized.
+  module Instrumentation
+    THREAD_KEY = :__search_engine_context__
+    THREAD_CORRELATION_KEY = :__search_engine_correlation_id__
+
+    CATALOG = {
+      'search_engine.search' => {
+        required: %i[collection],
+        optional: %i[
+          labels status duration_ms correlation_id params_preview http_status url_opts
+          error_class error_message selection_include_count selection_exclude_count
+          selection_nested_assoc_count preset_name preset_mode preset_pruned_keys
+          preset_pruned_keys_count preset_locked_domains_count curation_pinned_count
+          curation_hidden_count curation_has_override_tags curation_filter_flag
+          curation_conflict_type curation_conflict_count
+        ]
+      },
+      'search_engine.multi_search' => {
+        required: %i[searches_count],
+        optional: %i[labels status duration_ms correlation_id url_opts http_status]
+      },
+      'search_engine.compile' => {
+        required: %i[collection klass node_count source],
+        optional: %i[duration_ms]
+      },
+      'search_engine.grouping.compile' => {
+        required: %i[field],
+        optional: %i[collection limit missing_values duration_ms]
+      },
+      'search_engine.joins.compile' => {
+        required: %i[collection],
+        optional: %i[join_count assocs used_in include_len filter_len sort_len duration_ms has_joins]
+      },
+      'search_engine.selection.compile' => {
+        required: %i[],
+        optional: %i[include_count exclude_count nested_assoc_count]
+      },
+      'search_engine.preset.apply' => {
+        required: %i[preset_name],
+        optional: %i[preset_mode preset_pruned_keys preset_pruned_keys_count preset_locked_domains_count]
+      },
+      'search_engine.preset.conflict' => {
+        required: %i[type count],
+        optional: %i[limit]
+      },
+      'search_engine.curation.compile' => {
+        required: %i[],
+        optional: %i[pinned_count hidden_count has_override_tags filter_curated_hits]
+      },
+      'search_engine.curation.conflict' => {
+        required: %i[type count],
+        optional: %i[limit]
+      },
+      'search_engine.schema.diff' => {
+        required: %i[collection],
+        optional: %i[fields_changed_count added_count removed_count in_sync duration_ms]
+      },
+      'search_engine.schema.apply' => {
+        required: %i[collection],
+        optional: %i[physical_new new_physical alias_swapped dropped_count retention_deleted_count status duration_ms]
+      },
+      'search_engine.indexer.partition_start' => {
+        required: %i[collection into partition],
+        optional: %i[dispatch_mode job_id timestamp]
+      },
+      'search_engine.indexer.partition_finish' => {
+        required: %i[collection into partition batches_total docs_total success_total failed_total status],
+        optional: %i[duration_ms]
+      },
+      'search_engine.indexer.batch_import' => {
+        required: %i[collection into batch_index docs_count],
+        optional: %i[
+          success_count failure_count attempts http_status bytes_sent duration_ms
+          transient_retry error_sample
+        ]
+      },
+      'search_engine.indexer.delete_stale' => {
+        required: %i[collection into partition filter_hash status],
+        optional: %i[deleted_count duration_ms reason]
+      },
+      'search_engine.result.grouped_parsed' => { required: %i[collection groups_count], optional: %i[] },
+      'search_engine.joins.declared' => { required: %i[model name collection], optional: %i[] },
+      'search_engine.relation.group_by_updated' => {
+        required: %i[collection field],
+        optional: %i[limit missing_values]
+      },
+      'search_engine.facet.compile' => {
+        required: %i[],
+        optional: %i[collection fields_count queries_count max_facet_values sort_flags conflicts duration_ms]
+      },
+      'search_engine.highlight.compile' => {
+        required: %i[],
+        optional: %i[collection fields_count full_fields_count affix_tokens snippet_threshold tag_kind duration_ms]
+      },
+      'search_engine.synonyms.apply' => {
+        required: %i[],
+        optional: %i[collection use_synonyms use_stopwords source duration_ms]
+      },
+      'search_engine.geo.compile' => {
+        required: %i[],
+        optional: %i[collection filters_count shapes sort_mode radius_bucket duration_ms]
+      },
+      'search_engine.vector.compile' => {
+        required: %i[],
+        optional: %i[collection query_vector_present dims hybrid_weight ann_params_present duration_ms]
+      },
+      'search_engine.hits.limit' => {
+        required: %i[],
+        optional: %i[collection early_limit validate_max applied_strategy triggered total_hits duration_ms]
+      }
+    }.freeze
+
+    def self.catalog
+      CATALOG
+    end
+
+    # Emit an event with a shaped payload and yield a mutable context.
+    # Stamps correlation_id, status, error_class/error_message, and duration_ms.
+    # @param event [String]
+    # @param payload [Hash]
+    # @yieldparam ctx [Hash]
+    # @return [Object] block result when provided
+    def self.instrument(event, payload = {})
+      started = monotonic_ms
+      ctx = shape_payload(payload)
+      ctx[:correlation_id] ||= (Thread.current[THREAD_CORRELATION_KEY] ||= generate_correlation_id)
+
+      if defined?(ActiveSupport::Notifications)
+        result = nil
+        ActiveSupport::Notifications.instrument(event, ctx) do
+          result = yield(ctx) if block_given?
+          ctx[:status] = :ok unless ctx.key?(:status)
+          result
+        rescue StandardError => error
+          fill_error_context!(ctx, error)
+          raise
+        ensure
+          ctx[:duration_ms] = (monotonic_ms - started).round(1)
+        end
+        return result
+      end
+
+      # Fallback path when AS::N is unavailable
+      begin
+        result = yield(ctx) if block_given?
+        ctx[:status] = :ok unless ctx.key?(:status)
+        result
+      rescue StandardError => error
+        fill_error_context!(ctx, error)
+        raise
+      ensure
+        ctx[:duration_ms] = (monotonic_ms - started).round(1)
+      end
+    end
+
+    # Known events
+    #
+    # - "search_engine.joins.compile": emitted once per relation compile summarizing JOIN usage.
+    #   Payload keys (nil/empty omitted):
+    #   - :collection [String]
+    #   - :join_count [Integer]
+    #   - :assocs [Array<String>]
+    #   - :used_in [Hash{Symbol=>Array<String>}] include/filter/sort association usage
+    #   - :include_len [Integer]
+    #   - :filter_len [Integer]
+    #   - :sort_len [Integer]
+    #   - :duration_ms [Float]
+    #   - :has_joins [Boolean]
+    # - "search_engine.selection.compile": emitted once per relation compile summarizing field selection counts.
+    #   Payload keys:
+    #   - :include_count [Integer] total effective include fields (root + nested after precedence)
+    #   - :exclude_count [Integer] total excluded fields (root + nested)
+    #   - :nested_assoc_count [Integer] associations with any selection state (include or exclude)
+    # - "search_engine.preset.apply": emitted once per relation compile when a preset is present.
+    #   Payload keys (keys-only; values redacted elsewhere):
+    #   - :preset_name [String] effective namespaced preset
+    #   - :mode [Symbol] one of :merge, :only, :lock
+    #   - :locked_domains [Array<Symbol>] configured locked domains for :lock mode
+    #   - :pruned_keys [Array<Symbol>] keys removed by the chosen mode
+    # - "search_engine.curation.compile": emitted once per relation compile when curation state is present.
+    #   Payload keys:
+    #   - :pinned_count [Integer]
+    #   - :hidden_count [Integer]
+    #   - :has_override_tags [Boolean]
+    #   - :filter_curated_hits [true,false,nil]
+    # - "search_engine.curation.conflict": emitted when overlaps or limits are detected; at most once per compile.
+    #   Payload keys:
+    #   - :type [Symbol] one of :overlap, :limit_exceeded
+    #   - :count [Integer]
+    #   - :limit [Integer, optional] present when type==:limit_exceeded
+    #   # See also: https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/presets#observability
+    # Measure a block and attach duration_ms to payload.
+    # @param event [String]
+    # @param base_payload [Hash]
+    # @yieldreturn [Object] result of the block
+    # @return [Object]
+    def self.time(event, base_payload = {})
+      started = monotonic_ms
+      result = nil
+      instrument(event, base_payload.merge(started_at_ms: started)) do
+        result = yield(base_payload)
+      end
+      result
+
+      # NOTE: ActiveSupport::Notifications already attaches duration; callers
+      # should prefer event.duration on the subscriber side. We still provide
+      # started_at_ms in the payload for completeness.
+    end
+
+    # Set a shallow thread-local context for nested calls.
+    # @param ctx [Hash]
+    # @yield block executed with context applied
+    # @return [Object] result of the block
+    def self.with_context(ctx)
+      prev = Thread.current[THREAD_KEY]
+      Thread.current[THREAD_KEY] = (prev || {}).merge(ctx || {})
+      yield
+    ensure
+      Thread.current[THREAD_KEY] = prev
+    end
+
+    # Current shallow context hash.
+    # @return [Hash]
+    def self.context
+      (Thread.current[THREAD_KEY] || {}).dup
+    end
+
+    # Apply/propagate correlation id for the current execution (thread/fiber-local)
+    # and restore the previous value afterwards.
+    # @param id [String, nil]
+    # @yieldreturn [Object]
+    def self.with_correlation_id(id = nil)
+      previous = Thread.current[THREAD_CORRELATION_KEY]
+      Thread.current[THREAD_CORRELATION_KEY] = id || previous || generate_correlation_id
+      yield
+    ensure
+      Thread.current[THREAD_CORRELATION_KEY] = if previous.nil?
+                                                 nil
+                                               else
+                                                 previous
+                                               end
+    end
+
+    # @return [String, nil]
+    def self.current_correlation_id
+      Thread.current[THREAD_CORRELATION_KEY]
+    end
+
+    # Redact a payload-like structure (delegates to Observability)
+    def self.redact(value)
+      SearchEngine::Observability.redact(value)
+    end
+
+    # Monotonic clock in milliseconds.
+    # @return [Float]
+    def self.monotonic_ms
+      Process.clock_gettime(Process::CLOCK_MONOTONIC, :float_millisecond)
+    rescue StandardError
+      (Time.now.to_f * 1000.0)
+    end
+
+    # Internal: normalize and prune a payload hash.
+    def self.shape_payload(payload)
+      shaped = {}
+      (payload || {}).each do |k, v|
+        next if v.nil?
+
+        shaped[k.to_sym] = v
+      end
+      # Attach current context values without overriding explicit keys
+      ctx = context
+      ctx.each do |k, v|
+        shaped[k] = v unless shaped.key?(k)
+      end
+      shaped
+    end
+    private_class_method :shape_payload
+
+    def self.generate_correlation_id
+      require 'securerandom'
+      SecureRandom.urlsafe_base64(8)
+    end
+    private_class_method :generate_correlation_id
+
+    def self.fill_error_context!(ctx, error)
+      ctx[:status] = :error unless ctx.key?(:status)
+      ctx[:error_class] ||= error.class.name
+      ctx[:error_message] ||= SearchEngine::Observability.truncate_message(
+        error.message,
+        SearchEngine.config.observability&.max_message_length || 200
+      )
+    end
+  end
+end

data/lib/search_engine/joins/guard.rb

@@ -0,0 +1,202 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  module Joins
+    # Stateless guard for validating join usage across Relation and Parser.
+    #
+    # Public API is module functions that raise {SearchEngine::Errors::*}
+    # on invalid inputs; successful validations return nil.
+    module Guard
+      module_function
+
+      # Ensure the association exists on +klass+.
+      #
+      # @param klass [Class] model class
+      # @param assoc [Symbol, String]
+      # @raise [SearchEngine::Errors::InvalidJoin]
+      # @return [void]
+      def ensure_assoc_exists!(klass, assoc)
+        key = assoc.to_sym
+        cfg = safe_joins_config(klass)[key]
+        return if cfg
+
+        suggestions = suggest(key, safe_joins_config(klass).keys)
+        model_name = safe_class_name(klass)
+        msg = "association :#{key} is not declared on #{model_name} "
+        msg += "(declare with `join :#{key}, ...`)"
+        msg += suggestion_suffix(suggestions)
+        raise SearchEngine::Errors::InvalidJoin.new(
+          msg,
+          hint: (suggestions&.any? ? "Did you mean #{suggestions.map { |s| ":#{s}" }.join(', ')}?" : nil),
+          doc: 'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/joins#troubleshooting',
+          details: { assoc: key, known: safe_joins_config(klass).keys }
+        )
+      end
+
+      # Ensure the association config is complete (local_key and foreign_key present and non-blank).
+      #
+      # @param klass [Class]
+      # @param assoc [Symbol, String]
+      # @raise [SearchEngine::Errors::InvalidJoinConfig]
+      # @return [void]
+      def ensure_config_complete!(klass, assoc)
+        key = assoc.to_sym
+        cfg = safe_joins_config(klass)[key]
+        # If missing entirely, surface presence error via ensure_assoc_exists!
+        unless cfg
+          ensure_assoc_exists!(klass, key)
+          return
+        end
+
+        missing = []
+        missing << :local_key if blank?(cfg[:local_key])
+        missing << :foreign_key if blank?(cfg[:foreign_key])
+        return if missing.empty?
+
+        model_name = safe_class_name(klass)
+        msg = "join :#{key} on #{model_name} is missing "
+        msg += missing.map { |m| ":#{m}" }.join(' and ')
+        msg += ' (declare with `join '
+        msg += ":#{key}, collection: ..., local_key: ..., foreign_key: ...`)"
+        raise SearchEngine::Errors::InvalidJoinConfig.new(
+          msg,
+          hint: 'Declare local_key and foreign_key in join config.',
+          doc: 'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/joins#troubleshooting',
+          details: { assoc: key, missing: missing }
+        )
+      end
+
+      # Ensure the relation has applied the association via .joins(:assoc) before use.
+      #
+      # @param joins [Array<Symbol>, nil]
+      # @param assoc [Symbol, String]
+      # @param context [String] optional human-friendly action (e.g., "filtering/sorting")
+      # @raise [SearchEngine::Errors::JoinNotApplied]
+      # @return [void]
+      def ensure_join_applied!(joins, assoc, context: 'filtering/sorting')
+        key = assoc.to_sym
+        list = Array(joins)
+        return if list.include?(key)
+
+        raise SearchEngine::Errors::JoinNotApplied.new(
+          "Call .joins(:#{key}) before #{context} on #{key} fields",
+          doc: 'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/joins#troubleshooting',
+          details: { assoc: key, context: context }
+        )
+      end
+
+      # Validate that a joined field exists on the target collection when the
+      # target model is registered and exposes attributes.
+      #
+      # Best-effort: when registry or attributes are unavailable, no exception
+      # is raised (prefer low-noise behavior on missing metadata).
+      #
+      # @param assoc_cfg [Hash] normalized association config
+      # @param field [Symbol, String]
+      # @param source_klass [Class, nil] optional source model class for error messages
+      # @raise [SearchEngine::Errors::UnknownJoinField]
+      # @return [void]
+      def validate_joined_field!(assoc_cfg, field, source_klass: nil)
+        return if assoc_cfg.nil?
+
+        collection = assoc_cfg[:collection]
+        return if blank?(collection)
+
+        target_klass = begin
+          SearchEngine.collection_for(collection)
+        rescue StandardError
+          nil
+        end
+        return unless target_klass.respond_to?(:attributes)
+
+        known = Array(target_klass.attributes).map { |k, _| k.to_s }
+        known |= ['id'] # allow id implicitly on joined collections
+        return if known.empty?
+
+        fname = field.to_s
+        return if known.include?(fname)
+
+        suggestions = suggest(fname, known)
+        # Prefer the source model name for message clarity when provided
+        model_name = safe_class_name(source_klass || target_klass)
+        assoc_name = assoc_cfg[:name] || begin
+          # best effort: derive from collection
+          collection.to_s
+        end
+        msg = "UnknownJoinField: :#{fname} is not declared on association :#{assoc_name} for #{model_name}"
+        msg += suggestion_suffix(suggestions)
+        raise SearchEngine::Errors::UnknownJoinField.new(
+          msg,
+          details: { assoc: assoc_name, field: fname }
+        )
+      end
+
+      # Reject multi-hop paths like $authors.publisher.name
+      #
+      # @param path [String]
+      # @raise [SearchEngine::Errors::UnsupportedJoinNesting]
+      # @return [void]
+      def ensure_single_hop_path!(path)
+        return if path.to_s.count('.') <= 1
+
+        raise SearchEngine::Errors::UnsupportedJoinNesting.new(
+          'Only one join hop is supported: `$assoc.field`. Use a separate pipeline step to denormalize deeper paths.',
+          doc: 'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/joins#troubleshooting',
+          details: { path: path }
+        )
+      end
+
+      # --- internals -------------------------------------------------------
+
+      def suggest(input, known)
+        return [] if known.nil? || known.empty?
+
+        begin
+          require 'did_you_mean'
+          require 'did_you_mean/levenshtein'
+        rescue StandardError
+          return []
+        end
+
+        candidates = Array(known).map(&:to_s)
+        str = input.to_s
+        distances = candidates.each_with_object({}) do |cand, acc|
+          acc[cand] = DidYouMean::Levenshtein.distance(str, cand)
+        end
+        min = distances.values.min
+        return [] if min.nil? || min > 2
+
+        best = distances.select { |_, d| d == min }.keys.sort
+        best.first(3).map(&:to_sym)
+      end
+      private_class_method :suggest
+
+      def suggestion_suffix(suggestions)
+        return '' if suggestions.nil? || suggestions.empty?
+
+        tail = suggestions.map { |s| ":#{s}" }.join(', ')
+        " (did you mean #{tail}?)"
+      end
+      private_class_method :suggestion_suffix
+
+      def safe_joins_config(klass)
+        if klass.respond_to?(:joins_config)
+          klass.joins_config || {}
+        else
+          {}
+        end
+      end
+      private_class_method :safe_joins_config
+
+      def safe_class_name(klass)
+        klass.respond_to?(:name) && klass.name ? klass.name : klass.to_s
+      end
+      private_class_method :safe_class_name
+
+      def blank?(value)
+        value.nil? || value.to_s.strip.empty?
+      end
+      private_class_method :blank?
+    end
+  end
+end

data/lib/search_engine/joins/resolver.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  module Joins
+    # Key resolver for join associations.
+    #
+    # Determines { local_key, foreign_key } to use for an association, honoring
+    # explicitly provided keys and auto-inferring when absent. Auto-inference
+    # prefers a single shared attribute name between the base and target models
+    # that ends with `_id`. When ambiguity is detected, an error is raised with
+    # short suggestions.
+    module Resolver
+      module_function
+
+      # Resolve join keys for an association.
+      #
+      # @param base_klass [Class] the model on which the join is declared
+      # @param assoc_cfg [Hash] normalized join config (from joins_config)
+      # @return [Hash] { local_key: Symbol, foreign_key: Symbol }
+      # @raise [SearchEngine::Errors::InvalidJoinConfig]
+      def resolve_keys(base_klass, assoc_cfg)
+        local_key = assoc_cfg[:local_key]
+        foreign_key = assoc_cfg[:foreign_key]
+
+        if present?(local_key) && present?(foreign_key)
+          validate_key_known!(base_klass, local_key, side: :local)
+          validate_foreign_key_known!(assoc_cfg, foreign_key)
+          return { local_key: local_key.to_sym, foreign_key: foreign_key.to_sym }
+        end
+
+        # Infer when any of the keys is missing
+        base_attrs = safe_attributes(base_klass).keys.map(&:to_s)
+        target_klass = safe_target_klass(assoc_cfg[:collection])
+        target_attrs = safe_attributes(target_klass).keys.map(&:to_s)
+
+        shared = (base_attrs & target_attrs)
+        candidates = shared.select { |n| n.end_with?('_id') }
+        candidates = shared if candidates.empty?
+
+        raise_ambiguous_keys!(base_klass, assoc_cfg, candidates) if candidates.length != 1
+
+        key = candidates.first.to_sym
+        { local_key: key, foreign_key: key }
+      end
+
+      # --- internals --------------------------------------------------------
+
+      def validate_key_known!(klass, key, side:)
+        attrs = safe_attributes(klass)
+        return if key.to_sym == :id || attrs.key?(key.to_sym)
+
+        model_name = klass.respond_to?(:name) && klass.name ? klass.name : klass.to_s
+        raise SearchEngine::Errors::InvalidJoin.new(
+          "Unknown #{side} key :#{key} for #{model_name}. Declare it via `attribute :#{key}, ...`.",
+          doc: 'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/joins#troubleshooting',
+          details: { side: side, key: key, model: model_name }
+        )
+      end
+
+      def validate_foreign_key_known!(assoc_cfg, key)
+        tklass = safe_target_klass(assoc_cfg[:collection])
+        validate_key_known!(tklass, key, side: :foreign)
+      end
+
+      def safe_target_klass(collection_name)
+        SearchEngine.collection_for(collection_name)
+      rescue StandardError
+        nil
+      end
+
+      def safe_attributes(klass)
+        return {} unless klass.respond_to?(:attributes)
+
+        klass.attributes || {}
+      end
+
+      def raise_ambiguous_keys!(base_klass, assoc_cfg, candidates)
+        base_name = base_klass.respond_to?(:name) && base_klass.name ? base_klass.name : base_klass.to_s
+        assoc_name = assoc_cfg[:name] || assoc_cfg[:collection] || :unknown
+        sugg = candidates.map { |n| ":#{n}" }.join(', ')
+        msg = "Ambiguous join keys for :#{assoc_name} on #{base_name}. " \
+              "Could not infer a unique shared key. Candidates: #{sugg}"
+        raise SearchEngine::Errors::InvalidJoinConfig.new(
+          msg,
+          doc: 'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/joins#client-side-fallback',
+          details: { assoc: assoc_name, candidates: candidates }
+        )
+      end
+
+      def present?(v)
+        !(v.nil? || v.to_s.strip.empty?)
+      end
+    end
+  end
+end

data/lib/search_engine/logging/color.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  module Logging
+    # ANSI color helpers for CLI output.
+    #
+    # Applies colors only when $stdout is a TTY and NO_COLOR is not set.
+    # Intended for short substring coloring inside existing log lines.
+    #
+    # @since M8
+    module Color
+      module_function
+
+      # @param str [String]
+      # @param color [Symbol] one of :green, :yellow, :red
+      # @return [String]
+      def apply(str, color)
+        return str unless enabled?
+
+        code = case color.to_sym
+               when :green then 32
+               when :yellow then 33
+               when :red then 31
+               else 0
+               end
+        return str if code.zero?
+
+        "\e[#{code}m#{str}\e[0m"
+      end
+
+      # Apply bold styling to a string.
+      # @param str [String]
+      # @return [String]
+      def bold(str)
+        return str unless enabled?
+
+        "\e[1m#{str}\e[0m"
+      end
+
+      # Map indexation status to a color.
+      # @param status [#to_s]
+      # @return [Symbol] color name
+      def for_status(status)
+        case status.to_s
+        when 'ok' then :green
+        when 'failed' then :red
+        when 'partial' then :yellow
+        else :yellow
+        end
+      end
+
+      # Determine color for partition/status based on success/failure counts.
+      # @param failed_total [Integer] number of failed documents
+      # @param success_total [Integer] number of successful documents
+      # @return [Symbol] color name (:green, :yellow, or :red)
+      def for_partition_status(failed_total, success_total)
+        if failed_total.to_i.zero?
+          :green
+        elsif success_total.to_i.positive?
+          :yellow # partial success
+        else
+          :red # all failed
+        end
+      end
+
+      # @return [Boolean] whether coloring is active
+      def enabled?
+        return false if ENV['NO_COLOR']
+
+        begin
+          $stdout.isatty
+        rescue StandardError
+          false
+        end
+      end
+    end
+  end
+end