From fe4540d11c8e9b34d89101e37d73b7d0b5b02e6e Mon Sep 17 00:00:00 2001
From: Rushik Subba <srushik@gitlab.com>
Date: Thu, 16 Oct 2025 17:48:37 +0530
Subject: [PATCH 1/4] Add vulnerabilities risk score finder

Adds VulnerabilityElasticRiskScoresFinder
which calculates total risk scores for a given
group or a project.

Changelog: added
EE: true
---
 ...ulnerability_elastic_risk_scores_finder.rb | 41 ++++++++
 .../elastic/vulnerability_aggregations.rb     | 20 ++++
 .../elastic/vulnerability_query_builder.rb    |  3 +
 ee/spec/factories/vulnerabilities.rb          | 17 ++++
 ee/spec/factories/vulnerabilities/findings.rb |  6 ++
 ...ability_elastic_risk_scores_finder_spec.rb | 97 +++++++++++++++++++
 .../vulnerability_aggregations_spec.rb        | 34 +++++++
 .../vulnerability_query_builder_spec.rb       | 14 +++
 spec/support/finder_collection_allowlist.yml  |  1 +
 9 files changed, 233 insertions(+)
 create mode 100644 ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
 create mode 100644 ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb

diff --git a/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
new file mode 100644
index 00000000000000..c9407d0196c1d2
--- /dev/null
+++ b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module Security
+  class VulnerabilityElasticRiskScoresFinder < VulnerabilityElasticBaseFinder # rubocop:disable Search/NamespacedClass -- Keeping this in the Security namespace as it is used specifically by security classes
+    def initialize(vulnerable, params = {})
+      super(vulnerable, params.merge({ state: Vulnerability.active_states }))
+    end
+
+    def execute
+      options = search_params.merge({ risk_score: true })
+
+      query = ::Search::Elastic::VulnerabilityQueryBuilder.build(query: nil, options: options)
+      Gitlab::Search::Client.execute_search(query: query, options: es_search_options) do |es_response|
+        response = ::Search::Elastic::ResponseMapper.new(es_response)
+        calculate_risk_score(response)
+      end
+    end
+
+    private
+
+    def calculate_risk_score(response)
+      risk_scores_sum = response.aggregations.dig("risk_scores_sum", "value") || 0
+      created_at_sum = response.aggregations.dig("created_at_sum", "value") || 0
+      active_vulnerabilities_count = response.total_count || 0
+      active_vulnerabilities_count = 1 if active_vulnerabilities_count == 0
+
+      # division by 1000 because ES returns milli seconds where as Time.zone.now.to_i gives seconds
+      date_diff = (active_vulnerabilities_count * Time.zone.now.to_f) - (created_at_sum / 1000)
+      date_diff_in_months = date_diff / 86400 / 30
+
+      total_risk_score(risk_scores_sum, date_diff_in_months, active_vulnerabilities_count)
+    end
+
+    # Risk Score = Sum(Vulnerability Scores) + Sum(Vulnerability_age_in_month) x 0.005) / sqrt(num_vulnerabilities).
+    def total_risk_score(risk_scores_sum, age_in_months, active_vulnerabilities_count)
+      group_risk_score = (risk_scores_sum + (0.005 * age_in_months)) / Math.sqrt(active_vulnerabilities_count)
+
+      [1.0, group_risk_score.round(2)].min
+    end
+  end
+end
diff --git a/ee/lib/search/elastic/vulnerability_aggregations.rb b/ee/lib/search/elastic/vulnerability_aggregations.rb
index cc314dbb2a7637..f82ea03b258ad7 100644
--- a/ee/lib/search/elastic/vulnerability_aggregations.rb
+++ b/ee/lib/search/elastic/vulnerability_aggregations.rb
@@ -108,6 +108,26 @@ def by_identifiers_search(query_hash:, options:)
             }
           )
         end
+
+        def by_risk_score(query_hash:, options:)
+          return query_hash unless options[:risk_score]
+
+          query_hash.merge(
+            size: 0,
+            aggs: {
+              risk_scores_sum: {
+                sum: {
+                  field: "risk_score"
+                }
+              },
+              created_at_sum: {
+                sum: {
+                  field: "created_at"
+                }
+              }
+            }
+          )
+        end
       end
     end
   end
diff --git a/ee/lib/search/elastic/vulnerability_query_builder.rb b/ee/lib/search/elastic/vulnerability_query_builder.rb
index 81c703cee08273..1d77d8e507beec 100644
--- a/ee/lib/search/elastic/vulnerability_query_builder.rb
+++ b/ee/lib/search/elastic/vulnerability_query_builder.rb
@@ -65,6 +65,9 @@ def build # rubocop:disable Metrics/AbcSize -- need all the filters in one place
             query_hash: query_hash, options: options)
         end
 
+        query_hash = ::Search::Elastic::VulnerabilityAggregations.by_risk_score(
+          query_hash: query_hash, options: options)
+
         query_hash = ::Search::Elastic::Formats.source_fields(query_hash: query_hash, options: options)
         query_hash = ::Search::Elastic::Formats.page(query_hash: query_hash, options: options)
         query_hash = ::Search::Elastic::Formats.size(query_hash: query_hash, options: options)
diff --git a/ee/spec/factories/vulnerabilities.rb b/ee/spec/factories/vulnerabilities.rb
index 86efda00840f5c..a046e565a3f68a 100644
--- a/ee/spec/factories/vulnerabilities.rb
+++ b/ee/spec/factories/vulnerabilities.rb
@@ -155,6 +155,23 @@
       end
     end
 
+    trait :with_finding_risk_score do
+      after(:build) do |vulnerability|
+        finding = create(
+          :vulnerabilities_finding,
+          :identifier,
+          :with_finding_risk_score,
+          severity: vulnerability.severity,
+          description: vulnerability.description,
+          vulnerability: vulnerability,
+          report_type: vulnerability.report_type,
+          project: vulnerability.project
+        )
+
+        vulnerability.findings = [finding]
+      end
+    end
+
     trait :with_cluster_image_scanning_finding do
       transient do
         agent_id { '46357' }
diff --git a/ee/spec/factories/vulnerabilities/findings.rb b/ee/spec/factories/vulnerabilities/findings.rb
index d421bf95a7ebbd..823ee820a91808 100644
--- a/ee/spec/factories/vulnerabilities/findings.rb
+++ b/ee/spec/factories/vulnerabilities/findings.rb
@@ -804,6 +804,12 @@
       end
     end
 
+    trait :with_finding_risk_score do
+      after(:create) do |finding|
+        create(:vulnerability_finding_risk_score, finding: finding)
+      end
+    end
+
     ::Enums::Vulnerability.report_types.keys.each do |security_report_type|
       trait security_report_type do
         report_type { security_report_type }
diff --git a/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
new file mode 100644
index 00000000000000..67ef09e05b9d95
--- /dev/null
+++ b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Security::VulnerabilityElasticRiskScoresFinder, :elastic_delete_by_query, feature_category: :vulnerability_management do
+  let_it_be(:group) { create(:group) }
+  let_it_be(:project_1) { create(:project, namespace: group) }
+  let_it_be(:project_2) { create(:project, namespace: group) }
+
+  let_it_be(:low_vulnerability) do
+    create(:vulnerability, :with_finding_risk_score, :detected, severity: :low, project: project_1)
+  end
+
+  let_it_be(:medium_vulnerability) do
+    create(:vulnerability, :with_finding_risk_score, :detected, severity: :medium, project: project_2)
+  end
+
+  let_it_be(:high_vulnerability) do
+    create(:vulnerability, :with_finding_risk_score, :confirmed, severity: :high, project: project_1)
+  end
+
+  let_it_be(:critical_vulnerability) do
+    create(:vulnerability, :with_finding_risk_score, :confirmed, severity: :critical, project: project_2)
+  end
+
+  let_it_be(:inactive_vulnerability_1) do
+    create(:vulnerability, :with_finding_risk_score, :resolved, severity: :low, project: project_1)
+  end
+
+  let_it_be(:inactive_vulnerability_2) do
+    create(:vulnerability, :with_finding_risk_score, :dismissed, severity: :low, project: project_2)
+  end
+
+  let_it_be(:active_vulnerabilities) do
+    [low_vulnerability, medium_vulnerability, high_vulnerability, critical_vulnerability]
+  end
+
+  let_it_be(:inactive_vulnerabilities) { [inactive_vulnerability_1, inactive_vulnerability_2] }
+  let_it_be(:vulnerabilities) { active_vulnerabilities + inactive_vulnerabilities }
+
+  let(:params) { {} }
+
+  subject(:execute) { described_class.new(vulnerable, params).execute }
+
+  before do
+    stub_ee_application_setting(elasticsearch_search: true, elasticsearch_indexing: true)
+  end
+
+  describe '#execute' do
+    before do
+      Elastic::ProcessBookkeepingService.track!(*vulnerabilities)
+      ensure_elasticsearch_index!
+    end
+
+    context "when the vulnerable is a group" do
+      let(:vulnerable) { group }
+
+      it "calculates risk score for all the active vulnerabilities in the group" do
+        travel_to(60.days.from_now) do
+          risk_score = [1.0, total_risk_score(active_vulnerabilities)].min
+
+          is_expected.to eq(risk_score)
+        end
+      end
+    end
+
+    context "when the vulnerable is a project" do
+      let(:vulnerable) { project_1 }
+
+      it "calculates risk score for all the active vulnerabilities in the group" do
+        vulns = vulnerable.vulnerabilities.active
+        travel_to(60.days.from_now) do
+          risk_score = [1.0, total_risk_score(vulns)].min
+
+          is_expected.to eq(risk_score)
+        end
+      end
+    end
+  end
+
+  private
+
+  def total_risk_score(vulnerabilities)
+    risk_scores_sum = vulnerabilities.sum do |vulnerability|
+      Vulnerabilities::RiskScore.from_finding(vulnerability.findings.first).score
+    end
+    created_at_sum = vulnerabilities.sum { |vuln| vuln.created_at.to_f }
+    vulnerabilities_count = vulnerabilities.size
+
+    date_diff = (vulnerabilities_count * Time.zone.now.to_f) - created_at_sum
+    age_in_months = date_diff / 86400 / 30
+
+    group_risk_score = (risk_scores_sum + (0.005 * age_in_months)) / Math.sqrt(vulnerabilities_count)
+
+    [1.0, group_risk_score.round(2)].min
+  end
+end
diff --git a/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb b/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb
index 17821bf5734e99..7dd0a791a038aa 100644
--- a/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb
+++ b/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb
@@ -178,4 +178,38 @@
       })
     end
   end
+
+  describe '.by_risk_score' do
+    let(:options) { { risk_score: risk_score } }
+
+    context 'with risk_score set to false' do
+      let(:risk_score) { false }
+
+      it "does not modify the query" do
+        expect(described_class.by_risk_score(query_hash: query_hash, options: options)).to eq(query_hash)
+      end
+    end
+
+    context 'with risk_score set to true' do
+      let(:risk_score) { true }
+
+      it "adds size and aggs to query_hash" do
+        expect(described_class.by_risk_score(query_hash: query_hash, options: options)).to eq({
+          size: 0,
+          aggs: {
+            risk_scores_sum: {
+              sum: {
+                field: "risk_score"
+              }
+            },
+            created_at_sum: {
+              sum: {
+                field: "created_at"
+              }
+            }
+          }
+        })
+      end
+    end
+  end
 end
diff --git a/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb b/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb
index 28ebd7d905a31f..bcf43fe6b2c9ff 100644
--- a/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb
+++ b/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb
@@ -330,6 +330,20 @@
     end
   end
 
+  describe 'by_risk_score' do
+    let(:options) do
+      base_options.merge(
+        risk_score: true
+      )
+    end
+
+    it 'does add aggregation' do
+      query = build
+
+      expect(query).to include(aggs: { risk_scores_sum: be_a(Hash), created_at_sum: be_a(Hash) })
+    end
+  end
+
   describe 'validity_check' do
     let(:options) { base_options.merge(validity_check: [::Security::FindingTokenStatus.statuses.each_value.first]) }
 
diff --git a/spec/support/finder_collection_allowlist.yml b/spec/support/finder_collection_allowlist.yml
index b8ec1a7523852a..1b6f6a08f63a87 100644
--- a/spec/support/finder_collection_allowlist.yml
+++ b/spec/support/finder_collection_allowlist.yml
@@ -14,6 +14,7 @@
 - Security::VulnerabilityElasticIdentifierNamesFinder # Reason: The finder deals with Elasticsearch records and not DB records
 - Security::VulnerabilityElasticCountBySeverityFinder # Reason: The finder deals with Elasticsearch records and not DB records
 - Security::VulnerabilityElasticCountOverTimeFinder # Reason: The finder deals with Elasticsearch records and not DB records
+- Security::VulnerabilityElasticRiskScoresFinder # Reason: The finder deals with Elasticsearch records and not DB records
 - AuditEvents::CombinedAuditEventFinder # Reason: The finder combines result from 4 different tables and also returns cursor to next page
 - Security::AnalyzerGroupStatusFinder # Reason: To give accurate counts, return all analyzer types, even when there is no DB record
 - Ai::ClickHouseUsageEventsFinder # Reason: The finder's data is coming from Clickhouse and not Postgres, no ActiveRelation involved
-- 
GitLab


From a9588b5038798ed7ed101b5c102d5b6fbbcc332b Mon Sep 17 00:00:00 2001
From: Rushik Subba <srushik@gitlab.com>
Date: Tue, 21 Oct 2025 17:43:49 +0530
Subject: [PATCH 2/4] Address GitLabDuo review comments

---
 ...ulnerability_elastic_risk_scores_finder.rb | 14 ++++++++---
 ...ability_elastic_risk_scores_finder_spec.rb | 25 ++++++++++++++-----
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
index c9407d0196c1d2..8c1456b2e32cd4 100644
--- a/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
+++ b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
@@ -1,5 +1,13 @@
 # frozen_string_literal: true
 
+# Security::VulnerabilityElasticCountBySeverityFinder
+#
+# Used to calculate total risk score of all active vulnerabilities
+# in a project or a group.
+#
+# Arguments:
+#   see in ee/app/finders/security/vulnerability_elastic_base_finder.rb
+
 module Security
   class VulnerabilityElasticRiskScoresFinder < VulnerabilityElasticBaseFinder # rubocop:disable Search/NamespacedClass -- Keeping this in the Security namespace as it is used specifically by security classes
     def initialize(vulnerable, params = {})
@@ -22,10 +30,10 @@ def calculate_risk_score(response)
       risk_scores_sum = response.aggregations.dig("risk_scores_sum", "value") || 0
       created_at_sum = response.aggregations.dig("created_at_sum", "value") || 0
       active_vulnerabilities_count = response.total_count || 0
-      active_vulnerabilities_count = 1 if active_vulnerabilities_count == 0
+      return 0.0 if active_vulnerabilities_count == 0
 
       # division by 1000 because ES returns milli seconds where as Time.zone.now.to_i gives seconds
-      date_diff = (active_vulnerabilities_count * Time.zone.now.to_f) - (created_at_sum / 1000)
+      date_diff = (active_vulnerabilities_count * Time.zone.now.to_i) - (created_at_sum / 1000).to_i
       date_diff_in_months = date_diff / 86400 / 30
 
       total_risk_score(risk_scores_sum, date_diff_in_months, active_vulnerabilities_count)
@@ -35,7 +43,7 @@ def calculate_risk_score(response)
     def total_risk_score(risk_scores_sum, age_in_months, active_vulnerabilities_count)
       group_risk_score = (risk_scores_sum + (0.005 * age_in_months)) / Math.sqrt(active_vulnerabilities_count)
 
-      [1.0, group_risk_score.round(2)].min
+      [1.0, group_risk_score.to_f].min
     end
   end
 end
diff --git a/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
index 67ef09e05b9d95..3692a65150895d 100644
--- a/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
+++ b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
@@ -59,7 +59,7 @@
         travel_to(60.days.from_now) do
           risk_score = [1.0, total_risk_score(active_vulnerabilities)].min
 
-          is_expected.to eq(risk_score)
+          is_expected.to be_within(0.005).of(risk_score)
         end
       end
     end
@@ -67,12 +67,25 @@
     context "when the vulnerable is a project" do
       let(:vulnerable) { project_1 }
 
-      it "calculates risk score for all the active vulnerabilities in the group" do
+      it "calculates risk score for all the active vulnerabilities in the project" do
         vulns = vulnerable.vulnerabilities.active
         travel_to(60.days.from_now) do
           risk_score = [1.0, total_risk_score(vulns)].min
 
-          is_expected.to eq(risk_score)
+          is_expected.to be_within(0.005).of(risk_score)
+        end
+      end
+    end
+
+    context "when there are no active vulnerabilities" do
+      let_it_be(:vulnerable) { create(:project, namespace: group) }
+      let_it_be(:inactive_vuln) do
+        create(:vulnerability, :with_finding_risk_score, :dismissed, severity: :low, project: vulnerable)
+      end
+
+      it "returns risk score as 0.0" do
+        travel_to(60.days.from_now) do
+          is_expected.to eq(0)
         end
       end
     end
@@ -84,14 +97,14 @@ def total_risk_score(vulnerabilities)
     risk_scores_sum = vulnerabilities.sum do |vulnerability|
       Vulnerabilities::RiskScore.from_finding(vulnerability.findings.first).score
     end
-    created_at_sum = vulnerabilities.sum { |vuln| vuln.created_at.to_f }
+    created_at_sum = vulnerabilities.sum { |vuln| vuln.created_at.to_i }
     vulnerabilities_count = vulnerabilities.size
 
-    date_diff = (vulnerabilities_count * Time.zone.now.to_f) - created_at_sum
+    date_diff = (vulnerabilities_count * Time.zone.now.to_i) - created_at_sum
     age_in_months = date_diff / 86400 / 30
 
     group_risk_score = (risk_scores_sum + (0.005 * age_in_months)) / Math.sqrt(vulnerabilities_count)
 
-    [1.0, group_risk_score.round(2)].min
+    [1.0, group_risk_score.to_f].min
   end
 end
-- 
GitLab


From ab00ce6bb222d40d43b8e9fb0a00d5e071c55239 Mon Sep 17 00:00:00 2001
From: Rushik Subba <srushik@gitlab.com>
Date: Wed, 22 Oct 2025 15:48:10 +0530
Subject: [PATCH 3/4] Refactor aggregate risk score calculation

---
 ...ulnerability_elastic_risk_scores_finder.rb | 20 +++-------
 ee/lib/vulnerabilities/risk_score.rb          | 22 +++++++++++
 ...ability_elastic_risk_scores_finder_spec.rb | 13 +++----
 .../lib/vulnerabilities/risk_score_spec.rb    | 39 +++++++++++++++++++
 4 files changed, 73 insertions(+), 21 deletions(-)

diff --git a/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
index 8c1456b2e32cd4..3dade4f037673a 100644
--- a/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
+++ b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-# Security::VulnerabilityElasticCountBySeverityFinder
+# Security::VulnerabilityElasticRiskScoresFinder
 #
 # Used to calculate total risk score of all active vulnerabilities
 # in a project or a group.
@@ -30,20 +30,12 @@ def calculate_risk_score(response)
       risk_scores_sum = response.aggregations.dig("risk_scores_sum", "value") || 0
       created_at_sum = response.aggregations.dig("created_at_sum", "value") || 0
       active_vulnerabilities_count = response.total_count || 0
-      return 0.0 if active_vulnerabilities_count == 0
 
-      # division by 1000 because ES returns milli seconds where as Time.zone.now.to_i gives seconds
-      date_diff = (active_vulnerabilities_count * Time.zone.now.to_i) - (created_at_sum / 1000).to_i
-      date_diff_in_months = date_diff / 86400 / 30
-
-      total_risk_score(risk_scores_sum, date_diff_in_months, active_vulnerabilities_count)
-    end
-
-    # Risk Score = Sum(Vulnerability Scores) + Sum(Vulnerability_age_in_month) x 0.005) / sqrt(num_vulnerabilities).
-    def total_risk_score(risk_scores_sum, age_in_months, active_vulnerabilities_count)
-      group_risk_score = (risk_scores_sum + (0.005 * age_in_months)) / Math.sqrt(active_vulnerabilities_count)
-
-      [1.0, group_risk_score.to_f].min
+      Vulnerabilities::RiskScore.aggregate_score(
+        risk_scores_sum: risk_scores_sum,
+        created_at_sum: created_at_sum,
+        active_vulnerabilities_count: active_vulnerabilities_count
+      )
     end
   end
 end
diff --git a/ee/lib/vulnerabilities/risk_score.rb b/ee/lib/vulnerabilities/risk_score.rb
index 81b6bec5ec9d39..5b21b1b96c4eaf 100644
--- a/ee/lib/vulnerabilities/risk_score.rb
+++ b/ee/lib/vulnerabilities/risk_score.rb
@@ -2,6 +2,28 @@
 
 module Vulnerabilities
   class RiskScore
+    # Calculate aggregate risk score for scoped vulnerabilities
+    # This can be for a group or a project or any set of vulnerabilities.
+    # Risk Score = (Sum(Vulnerability Scores) + Sum(Vulnerability_age_in_month) x 0.005) / sqrt(num_vulnerabilities)
+    #
+    # Arguments -
+    #
+    #   risk_score_sum: sum of all vulnerability finding risk scores
+    #   created_at_sum: sum of epoch values for created_at, in milliseconds
+    #   active_vulnerabilities_count: number of active vulnerabilities
+    def self.aggregate_score(risk_scores_sum:, created_at_sum:, active_vulnerabilities_count:)
+      return 0.0 if active_vulnerabilities_count == 0
+
+      # division by 1000 because created_at_sum is in milli seconds where as Time.zone.now.to_f gives seconds
+      date_diff = (active_vulnerabilities_count * Time.zone.now.to_f) - (created_at_sum / 1000)
+      date_diff_in_months = date_diff / 86400 / 30
+
+      age_factor = 0.005 * date_diff_in_months
+      total_risk_score = ((risk_scores_sum + age_factor) / Math.sqrt(active_vulnerabilities_count)).round(4)
+
+      [1.0, total_risk_score.to_f].min
+    end
+
     attr_reader :severity, :epss_score, :is_known_exploit
 
     def self.from_finding(finding)
diff --git a/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
index 3692a65150895d..032b3a71f74648 100644
--- a/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
+++ b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
@@ -97,14 +97,13 @@ def total_risk_score(vulnerabilities)
     risk_scores_sum = vulnerabilities.sum do |vulnerability|
       Vulnerabilities::RiskScore.from_finding(vulnerability.findings.first).score
     end
-    created_at_sum = vulnerabilities.sum { |vuln| vuln.created_at.to_i }
+    created_at_sum = vulnerabilities.sum { |vuln| vuln.created_at.to_f * 1000 }
     vulnerabilities_count = vulnerabilities.size
 
-    date_diff = (vulnerabilities_count * Time.zone.now.to_i) - created_at_sum
-    age_in_months = date_diff / 86400 / 30
-
-    group_risk_score = (risk_scores_sum + (0.005 * age_in_months)) / Math.sqrt(vulnerabilities_count)
-
-    [1.0, group_risk_score.to_f].min
+    Vulnerabilities::RiskScore.aggregate_score(
+      risk_scores_sum: risk_scores_sum,
+      created_at_sum: created_at_sum,
+      active_vulnerabilities_count: vulnerabilities_count
+    )
   end
 end
diff --git a/ee/spec/lib/vulnerabilities/risk_score_spec.rb b/ee/spec/lib/vulnerabilities/risk_score_spec.rb
index d7b459dd5d1b98..e6ad99edb6d3a1 100644
--- a/ee/spec/lib/vulnerabilities/risk_score_spec.rb
+++ b/ee/spec/lib/vulnerabilities/risk_score_spec.rb
@@ -102,4 +102,43 @@
       end
     end
   end
+
+  describe ".aggregate_score" do
+    subject(:aggregate_score) do
+      described_class.aggregate_score(
+        risk_scores_sum: risk_scores_sum,
+        created_at_sum: created_at_sum,
+        active_vulnerabilities_count: active_vulnerabilities_count
+      )
+    end
+
+    let(:risk_scores_sum) { 1.5 }
+    let(:active_vulnerabilities_count) { 4 }
+    let(:created_at_sum) { active_vulnerabilities_count * 30.days.ago.to_f * 1000 }
+
+    context "when there are no active vulnerabilities" do
+      let(:active_vulnerabilities_count) { 0 }
+
+      it "returns 0.0" do
+        expect(aggregate_score).to eq(0.0)
+      end
+    end
+
+    context "when there are active vulnerabilities" do
+      it "calculates the aggregate risk score with age factor" do
+        expect(aggregate_score).to be > 0.0
+        expect(aggregate_score).to be <= 1.0
+
+        expect(aggregate_score).to be_within(0.005).of(0.76)
+      end
+    end
+
+    context "with very high risk scores" do
+      let(:risk_scores_sum) { 10.0 }
+
+      it "caps the result at 1.0" do
+        expect(aggregate_score).to eq(1.0)
+      end
+    end
+  end
 end
-- 
GitLab


From d3aa08f27daae29ec5ab8cde9b75aea586a1fca9 Mon Sep 17 00:00:00 2001
From: Rushik Subba <srushik@gitlab.com>
Date: Mon, 27 Oct 2025 17:50:00 +0530
Subject: [PATCH 4/4] Address review comments and refine code

1. Moves agg risk score calculation to seperate svc
2. Adds guard to risk score ES aggregation
3. Seperates risk score and created at ES aggregations
---
 ...ulnerability_elastic_risk_scores_finder.rb |  7 ++-
 .../elastic/vulnerability_aggregations.rb     | 18 ++++++--
 .../elastic/vulnerability_query_builder.rb    |  9 +++-
 .../vulnerabilities/aggregate_risk_score.rb   | 27 ++++++++++++
 ee/lib/vulnerabilities/risk_score.rb          | 22 ----------
 ...ability_elastic_risk_scores_finder_spec.rb |  2 +-
 .../vulnerability_aggregations_spec.rb        | 38 +++++++++++++---
 .../vulnerability_query_builder_spec.rb       | 38 ++++++++++++++--
 .../aggregate_risk_score_spec.rb              | 44 +++++++++++++++++++
 .../lib/vulnerabilities/risk_score_spec.rb    | 39 ----------------
 10 files changed, 164 insertions(+), 80 deletions(-)
 create mode 100644 ee/lib/vulnerabilities/aggregate_risk_score.rb
 create mode 100644 ee/spec/lib/vulnerabilities/aggregate_risk_score_spec.rb

diff --git a/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
index 3dade4f037673a..d41930b283701b 100644
--- a/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
+++ b/ee/app/finders/security/vulnerability_elastic_risk_scores_finder.rb
@@ -15,7 +15,10 @@ def initialize(vulnerable, params = {})
     end
 
     def execute
-      options = search_params.merge({ risk_score: true })
+      options = search_params.merge({
+        risk_score_sum: true,
+        created_at_sum: true
+      })
 
       query = ::Search::Elastic::VulnerabilityQueryBuilder.build(query: nil, options: options)
       Gitlab::Search::Client.execute_search(query: query, options: es_search_options) do |es_response|
@@ -31,7 +34,7 @@ def calculate_risk_score(response)
       created_at_sum = response.aggregations.dig("created_at_sum", "value") || 0
       active_vulnerabilities_count = response.total_count || 0
 
-      Vulnerabilities::RiskScore.aggregate_score(
+      Vulnerabilities::AggregateRiskScore.score(
         risk_scores_sum: risk_scores_sum,
         created_at_sum: created_at_sum,
         active_vulnerabilities_count: active_vulnerabilities_count
diff --git a/ee/lib/search/elastic/vulnerability_aggregations.rb b/ee/lib/search/elastic/vulnerability_aggregations.rb
index f82ea03b258ad7..fb8e99092f3726 100644
--- a/ee/lib/search/elastic/vulnerability_aggregations.rb
+++ b/ee/lib/search/elastic/vulnerability_aggregations.rb
@@ -109,17 +109,27 @@ def by_identifiers_search(query_hash:, options:)
           )
         end
 
-        def by_risk_score(query_hash:, options:)
-          return query_hash unless options[:risk_score]
+        def by_risk_score_sum(query_hash:, options:)
+          return query_hash unless options[:risk_score_sum]
 
-          query_hash.merge(
+          query_hash.deep_merge(
             size: 0,
             aggs: {
               risk_scores_sum: {
                 sum: {
                   field: "risk_score"
                 }
-              },
+              }
+            }
+          )
+        end
+
+        def by_created_at_sum(query_hash:, options:)
+          return query_hash unless options[:created_at_sum]
+
+          query_hash.deep_merge(
+            size: 0,
+            aggs: {
               created_at_sum: {
                 sum: {
                   field: "created_at"
diff --git a/ee/lib/search/elastic/vulnerability_query_builder.rb b/ee/lib/search/elastic/vulnerability_query_builder.rb
index 1d77d8e507beec..16144f305dea87 100644
--- a/ee/lib/search/elastic/vulnerability_query_builder.rb
+++ b/ee/lib/search/elastic/vulnerability_query_builder.rb
@@ -65,8 +65,13 @@ def build # rubocop:disable Metrics/AbcSize -- need all the filters in one place
             query_hash: query_hash, options: options)
         end
 
-        query_hash = ::Search::Elastic::VulnerabilityAggregations.by_risk_score(
-          query_hash: query_hash, options: options)
+        if ::Elastic::DataMigrationService.migration_has_finished?(:add_risk_score_field_to_vulnerability)
+          query_hash = ::Search::Elastic::VulnerabilityAggregations.by_risk_score_sum(query_hash: query_hash,
+            options: options)
+        end
+
+        query_hash = ::Search::Elastic::VulnerabilityAggregations.by_created_at_sum(query_hash: query_hash,
+          options: options)
 
         query_hash = ::Search::Elastic::Formats.source_fields(query_hash: query_hash, options: options)
         query_hash = ::Search::Elastic::Formats.page(query_hash: query_hash, options: options)
diff --git a/ee/lib/vulnerabilities/aggregate_risk_score.rb b/ee/lib/vulnerabilities/aggregate_risk_score.rb
new file mode 100644
index 00000000000000..84157e721253bc
--- /dev/null
+++ b/ee/lib/vulnerabilities/aggregate_risk_score.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Vulnerabilities
+  class AggregateRiskScore
+    # Calculate aggregate risk score for scoped vulnerabilities
+    # This can be for a group or a project or any set of vulnerabilities.
+    # Risk Score = (Sum(Vulnerability Scores) + Sum(Vulnerability_age_in_month) x 0.005) / sqrt(num_vulnerabilities)
+    #
+    # Arguments -
+    #
+    #   risk_score_sum: sum of all vulnerability finding risk scores
+    #   created_at_sum: sum of epoch values for created_at, in milliseconds
+    #   active_vulnerabilities_count: number of active vulnerabilities
+    def self.score(risk_scores_sum:, created_at_sum:, active_vulnerabilities_count:)
+      return 0.0 if active_vulnerabilities_count == 0
+
+      # division by 1000 because created_at_sum is in milli seconds where as Time.zone.now.to_f gives seconds
+      date_diff = (active_vulnerabilities_count * Time.zone.now.to_f) - (created_at_sum / 1000)
+      date_diff_in_months = date_diff / 86400 / 30
+
+      age_factor = 0.005 * date_diff_in_months
+      total_risk_score = ((risk_scores_sum + age_factor) / Math.sqrt(active_vulnerabilities_count)).round(4)
+
+      [1.0, total_risk_score.to_f].min
+    end
+  end
+end
diff --git a/ee/lib/vulnerabilities/risk_score.rb b/ee/lib/vulnerabilities/risk_score.rb
index 5b21b1b96c4eaf..81b6bec5ec9d39 100644
--- a/ee/lib/vulnerabilities/risk_score.rb
+++ b/ee/lib/vulnerabilities/risk_score.rb
@@ -2,28 +2,6 @@
 
 module Vulnerabilities
   class RiskScore
-    # Calculate aggregate risk score for scoped vulnerabilities
-    # This can be for a group or a project or any set of vulnerabilities.
-    # Risk Score = (Sum(Vulnerability Scores) + Sum(Vulnerability_age_in_month) x 0.005) / sqrt(num_vulnerabilities)
-    #
-    # Arguments -
-    #
-    #   risk_score_sum: sum of all vulnerability finding risk scores
-    #   created_at_sum: sum of epoch values for created_at, in milliseconds
-    #   active_vulnerabilities_count: number of active vulnerabilities
-    def self.aggregate_score(risk_scores_sum:, created_at_sum:, active_vulnerabilities_count:)
-      return 0.0 if active_vulnerabilities_count == 0
-
-      # division by 1000 because created_at_sum is in milli seconds where as Time.zone.now.to_f gives seconds
-      date_diff = (active_vulnerabilities_count * Time.zone.now.to_f) - (created_at_sum / 1000)
-      date_diff_in_months = date_diff / 86400 / 30
-
-      age_factor = 0.005 * date_diff_in_months
-      total_risk_score = ((risk_scores_sum + age_factor) / Math.sqrt(active_vulnerabilities_count)).round(4)
-
-      [1.0, total_risk_score.to_f].min
-    end
-
     attr_reader :severity, :epss_score, :is_known_exploit
 
     def self.from_finding(finding)
diff --git a/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
index 032b3a71f74648..134518c210cc1c 100644
--- a/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
+++ b/ee/spec/finders/security/vulnerability_elastic_risk_scores_finder_spec.rb
@@ -100,7 +100,7 @@ def total_risk_score(vulnerabilities)
     created_at_sum = vulnerabilities.sum { |vuln| vuln.created_at.to_f * 1000 }
     vulnerabilities_count = vulnerabilities.size
 
-    Vulnerabilities::RiskScore.aggregate_score(
+    Vulnerabilities::AggregateRiskScore.score(
       risk_scores_sum: risk_scores_sum,
       created_at_sum: created_at_sum,
       active_vulnerabilities_count: vulnerabilities_count
diff --git a/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb b/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb
index 7dd0a791a038aa..167a1d1499d124 100644
--- a/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb
+++ b/ee/spec/lib/search/elastic/vulnerability_aggregations_spec.rb
@@ -179,29 +179,53 @@
     end
   end
 
-  describe '.by_risk_score' do
-    let(:options) { { risk_score: risk_score } }
+  describe '.by_risk_score_sum' do
+    let(:options) { { risk_score_sum: risk_score_sum } }
 
     context 'with risk_score set to false' do
-      let(:risk_score) { false }
+      let(:risk_score_sum) { false }
 
       it "does not modify the query" do
-        expect(described_class.by_risk_score(query_hash: query_hash, options: options)).to eq(query_hash)
+        expect(described_class.by_risk_score_sum(query_hash: query_hash, options: options)).to eq(query_hash)
       end
     end
 
     context 'with risk_score set to true' do
-      let(:risk_score) { true }
+      let(:risk_score_sum) { true }
 
       it "adds size and aggs to query_hash" do
-        expect(described_class.by_risk_score(query_hash: query_hash, options: options)).to eq({
+        expect(described_class.by_risk_score_sum(query_hash: query_hash, options: options)).to eq({
           size: 0,
           aggs: {
             risk_scores_sum: {
               sum: {
                 field: "risk_score"
               }
-            },
+            }
+          }
+        })
+      end
+    end
+  end
+
+  describe '.by_created_at_sum' do
+    let(:options) { { created_at_sum: created_at_sum } }
+
+    context 'with created_at set to false' do
+      let(:created_at_sum) { false }
+
+      it "does not modify the query" do
+        expect(described_class.by_created_at_sum(query_hash: query_hash, options: options)).to eq(query_hash)
+      end
+    end
+
+    context 'with created_at set to true' do
+      let(:created_at_sum) { true }
+
+      it "adds size and aggs to query_hash" do
+        expect(described_class.by_created_at_sum(query_hash: query_hash, options: options)).to eq({
+          size: 0,
+          aggs: {
             created_at_sum: {
               sum: {
                 field: "created_at"
diff --git a/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb b/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb
index bcf43fe6b2c9ff..305c1db6a43eae 100644
--- a/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb
+++ b/ee/spec/lib/search/elastic/vulnerability_query_builder_spec.rb
@@ -330,17 +330,49 @@
     end
   end
 
-  describe 'by_risk_score' do
+  describe 'by_risk_score_sum' do
     let(:options) do
       base_options.merge(
-        risk_score: true
+        risk_score_sum: true
+      )
+    end
+
+    context 'when the es migration is completed' do
+      before do
+        set_elasticsearch_migration_to(:add_risk_score_field_to_vulnerability)
+      end
+
+      it 'does add aggregation' do
+        query = build
+
+        expect(query).to include(aggs: { risk_scores_sum: be_a(Hash) })
+      end
+    end
+
+    context 'when the es backfill migration is not completed' do
+      before do
+        set_elasticsearch_migration_to(:add_risk_score_field_to_vulnerability, including: false)
+      end
+
+      it 'does not add aggregation' do
+        query = build
+
+        expect(query).not_to include(aggs: { risk_scores_sum: be_a(Hash) })
+      end
+    end
+  end
+
+  describe 'by_created_at_sum' do
+    let(:options) do
+      base_options.merge(
+        created_at_sum: true
       )
     end
 
     it 'does add aggregation' do
       query = build
 
-      expect(query).to include(aggs: { risk_scores_sum: be_a(Hash), created_at_sum: be_a(Hash) })
+      expect(query).to include(aggs: { created_at_sum: be_a(Hash) })
     end
   end
 
diff --git a/ee/spec/lib/vulnerabilities/aggregate_risk_score_spec.rb b/ee/spec/lib/vulnerabilities/aggregate_risk_score_spec.rb
new file mode 100644
index 00000000000000..86e6aa9d905011
--- /dev/null
+++ b/ee/spec/lib/vulnerabilities/aggregate_risk_score_spec.rb
@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Vulnerabilities::AggregateRiskScore, feature_category: :vulnerability_management do
+  describe ".score" do
+    subject(:score) do
+      described_class.score(
+        risk_scores_sum: risk_scores_sum,
+        created_at_sum: created_at_sum,
+        active_vulnerabilities_count: active_vulnerabilities_count
+      )
+    end
+
+    let(:risk_scores_sum) { 1.5 }
+    let(:active_vulnerabilities_count) { 4 }
+    let(:created_at_sum) { active_vulnerabilities_count * 30.days.ago.to_f * 1000 }
+
+    context "when there are no active vulnerabilities" do
+      let(:active_vulnerabilities_count) { 0 }
+
+      it "returns 0.0" do
+        expect(score).to eq(0.0)
+      end
+    end
+
+    context "when there are active vulnerabilities" do
+      it "calculates the aggregate risk score with age factor" do
+        expect(score).to be > 0.0
+        expect(score).to be <= 1.0
+
+        expect(score).to be_within(0.005).of(0.76)
+      end
+    end
+
+    context "with very high risk scores" do
+      let(:risk_scores_sum) { 10.0 }
+
+      it "caps the result at 1.0" do
+        expect(score).to eq(1.0)
+      end
+    end
+  end
+end
diff --git a/ee/spec/lib/vulnerabilities/risk_score_spec.rb b/ee/spec/lib/vulnerabilities/risk_score_spec.rb
index e6ad99edb6d3a1..d7b459dd5d1b98 100644
--- a/ee/spec/lib/vulnerabilities/risk_score_spec.rb
+++ b/ee/spec/lib/vulnerabilities/risk_score_spec.rb
@@ -102,43 +102,4 @@
       end
     end
   end
-
-  describe ".aggregate_score" do
-    subject(:aggregate_score) do
-      described_class.aggregate_score(
-        risk_scores_sum: risk_scores_sum,
-        created_at_sum: created_at_sum,
-        active_vulnerabilities_count: active_vulnerabilities_count
-      )
-    end
-
-    let(:risk_scores_sum) { 1.5 }
-    let(:active_vulnerabilities_count) { 4 }
-    let(:created_at_sum) { active_vulnerabilities_count * 30.days.ago.to_f * 1000 }
-
-    context "when there are no active vulnerabilities" do
-      let(:active_vulnerabilities_count) { 0 }
-
-      it "returns 0.0" do
-        expect(aggregate_score).to eq(0.0)
-      end
-    end
-
-    context "when there are active vulnerabilities" do
-      it "calculates the aggregate risk score with age factor" do
-        expect(aggregate_score).to be > 0.0
-        expect(aggregate_score).to be <= 1.0
-
-        expect(aggregate_score).to be_within(0.005).of(0.76)
-      end
-    end
-
-    context "with very high risk scores" do
-      let(:risk_scores_sum) { 10.0 }
-
-      it "caps the result at 1.0" do
-        expect(aggregate_score).to eq(1.0)
-      end
-    end
-  end
 end
-- 
GitLab