diff --git a/doc/api/graphql/reference/index.md b/doc/api/graphql/reference/index.md index 72cae45176e7ead5c651c2bcae4efef31f847cb7..920ae62b51855792a9a427edcaa84a869c0c80f8 100644 --- a/doc/api/graphql/reference/index.md +++ b/doc/api/graphql/reference/index.md @@ -960,6 +960,7 @@ Input type: `AiActionInput` | ---- | ---- | ----------- | | <a id="mutationaiactionclientmutationid"></a>`clientMutationId` | [`String`](#string) | A unique identifier for the client performing the mutation. | | <a id="mutationaiactionexplaincode"></a>`explainCode` | [`AiExplainCodeInput`](#aiexplaincodeinput) | Input for explain_code AI action. | +| <a id="mutationaiactionexplainvulnerability"></a>`explainVulnerability` | [`AiExplainVulnerabilityInput`](#aiexplainvulnerabilityinput) | Input for explain_vulnerability AI action. | | <a id="mutationaiactionsummarizecomments"></a>`summarizeComments` | [`AiSummarizeCommentsInput`](#aisummarizecommentsinput) | Input for summarize_comments AI action. | #### Fields @@ -26500,6 +26501,14 @@ see the associated mutation type above. | <a id="aiexplaincodemessageinputcontent"></a>`content` | [`String!`](#string) | Content of the message. | | <a id="aiexplaincodemessageinputrole"></a>`role` | [`String!`](#string) | Role of the message (system, user, assistant). | +### `AiExplainVulnerabilityInput` + +#### Arguments + +| Name | Type | Description | +| ---- | ---- | ----------- | +| <a id="aiexplainvulnerabilityinputresourceid"></a>`resourceId` | [`AiModelID!`](#aimodelid) | GID of the resource to mutate. | + ### `AiSummarizeCommentsInput` #### Arguments diff --git a/ee/app/graphql/types/ai/explain_vulnerability_input_type.rb b/ee/app/graphql/types/ai/explain_vulnerability_input_type.rb new file mode 100644 index 0000000000000000000000000000000000000000..82d00ac1e91018a37004507d9c9be3633091eeb0 --- /dev/null +++ b/ee/app/graphql/types/ai/explain_vulnerability_input_type.rb @@ -0,0 +1,9 @@ +# frozen_string_literal: true + +module Types + module Ai + class ExplainVulnerabilityInputType < BaseMethodInputType + graphql_name 'AiExplainVulnerabilityInput' + end + end +end diff --git a/ee/app/models/ee/vulnerability.rb b/ee/app/models/ee/vulnerability.rb index 090b6607a1983cbb35028829a00fc451c2580abe..ab6e4baf706bdc88b2e2f49fa3acce8e3ed84448 100644 --- a/ee/app/models/ee/vulnerability.rb +++ b/ee/app/models/ee/vulnerability.rb @@ -6,6 +6,7 @@ module Vulnerability extend ActiveSupport::Concern prepended do + include ::Ai::Model include ::CacheMarkdownField include ::Redactable include ::StripAttribute @@ -227,6 +228,11 @@ def latest_state_transition @latest_state_transition ||= state_transitions.last end + def send_to_ai? + ::Feature.enabled?(:explain_vulnerability, project) && + finding&.ai_explainable? + end + private def note_as_string(time, username, status, text) diff --git a/ee/app/models/vulnerabilities/finding.rb b/ee/app/models/vulnerabilities/finding.rb index e7b9c3f8358a8f329621eccf899a9685b43d4d88..29806f096ef40405ae198c898c160d6b2b150572 100644 --- a/ee/app/models/vulnerabilities/finding.rb +++ b/ee/app/models/vulnerabilities/finding.rb @@ -151,6 +151,16 @@ def state end end + def source_code? + source_code.present? + end + + def vulnerable_code(lines: (start_line..end_line)) + strong_memoize_with(:vulnerable_code, lines) do + source_code.lines[lines]&.join + end + end + def self.related_dismissal_feedback Feedback.where('vulnerability_occurrences.uuid::uuid = vulnerability_feedback.finding_uuid') .for_dismissal @@ -400,11 +410,35 @@ def last_finding_pipeline finding_pipelines.last&.pipeline end + def ai_explainable? + file.present? && location["start_line"].present? + end + protected def primary_identifier_fingerprint identifiers.first&.fingerprint end + + private + + def start_line + [location["start_line"].to_i - 1, 0].max + end + + def end_line + return -1 if location["end_line"].blank? + + [location["end_line"].to_i - 1, start_line].max + end + + def source_code + return "" unless file.present? + + blob = project.repository.blob_at(pipeline_branch, file) + blob.present? ? blob.data : "" + end + strong_memoize_attr :source_code end end diff --git a/ee/app/policies/vulnerability_policy.rb b/ee/app/policies/vulnerability_policy.rb index 92eccf98cb27ec5915fda5bb8d8e99c443f27f26..0205151f127e65a9d6d29f6423ea09b1f51ec11b 100644 --- a/ee/app/policies/vulnerability_policy.rb +++ b/ee/app/policies/vulnerability_policy.rb @@ -7,4 +7,5 @@ class VulnerabilityPolicy < BasePolicy # It would not be safe to prevent :create_note in project policy, # since note permissions are shared, and this can have ripple effect on other parts. rule { ~can?(:read_security_resource) }.prevent :create_note + rule { can?(:read_security_resource) }.enable :read_vulnerability end diff --git a/ee/app/services/llm/execute_method_service.rb b/ee/app/services/llm/execute_method_service.rb index 0de44f00e9aab0f97ff2987c37a6a3846243c13c..0f07da60429fcac210415b0576876a17b4061947 100644 --- a/ee/app/services/llm/execute_method_service.rb +++ b/ee/app/services/llm/execute_method_service.rb @@ -5,6 +5,7 @@ class ExecuteMethodService < BaseService # This list of methods will expand as we add more methods to support. # Could also be abstracted to another class specific to find the appropriate method service. METHODS = { + explain_vulnerability: ::Llm::ExplainVulnerabilityService, summarize_comments: Llm::GenerateSummaryService, explain_code: Llm::ExplainCodeService }.freeze diff --git a/ee/app/services/llm/explain_vulnerability_service.rb b/ee/app/services/llm/explain_vulnerability_service.rb new file mode 100644 index 0000000000000000000000000000000000000000..dafc45d9923a115ea15f5c9c150fd75d3a7246a9 --- /dev/null +++ b/ee/app/services/llm/explain_vulnerability_service.rb @@ -0,0 +1,18 @@ +# frozen_string_literal: true + +module Llm + class ExplainVulnerabilityService < BaseService + private + + def perform + ::Llm::CompletionWorker.perform_async( + user.id, + resource.id, + resource.class.name, + :explain_vulnerability, + options + ) + success + end + end +end diff --git a/ee/lib/gitlab/llm/open_ai/completions/explain_vulnerability.rb b/ee/lib/gitlab/llm/open_ai/completions/explain_vulnerability.rb new file mode 100644 index 0000000000000000000000000000000000000000..c01bd5980435809a691fb6c20a50d855005eb985 --- /dev/null +++ b/ee/lib/gitlab/llm/open_ai/completions/explain_vulnerability.rb @@ -0,0 +1,42 @@ +# frozen_string_literal: true + +module Gitlab + module Llm + module OpenAi + module Completions + class ExplainVulnerability + DEFAULT_ERROR = '{"error": "An unexpected error has occurred."}' + + def initialize(template_class) + @template_class = template_class + end + + def execute(user, vulnerability, _options) + template = template_class.new(vulnerability) + response = response_for(user, template) + + ::Gitlab::Llm::OpenAi::ResponseService + .new(user, vulnerability, response, options: {}) + .execute(Gitlab::Llm::OpenAi::ResponseModifiers::Chat.new) + rescue StandardError => error + Gitlab::ErrorTracking.track_exception(error) + + ::Gitlab::Llm::OpenAi::ResponseService + .new(user, vulnerability, DEFAULT_ERROR, options: {}) + .execute + end + + private + + attr_reader :template_class + + def response_for(user, template) + Gitlab::Llm::OpenAi::Client + .new(user) + .chat(content: template.to_prompt, **template.options) + end + end + end + end + end +end diff --git a/ee/lib/gitlab/llm/open_ai/completions/factory.rb b/ee/lib/gitlab/llm/open_ai/completions/factory.rb index c1ddc0bd1f68c37a89eca2cb2e32ac83abfedf23..b59f7dcb332e693bc9f3b014a11e3cc76612b769 100644 --- a/ee/lib/gitlab/llm/open_ai/completions/factory.rb +++ b/ee/lib/gitlab/llm/open_ai/completions/factory.rb @@ -6,6 +6,10 @@ module OpenAi module Completions class Factory COMPLETIONS = { + explain_vulnerability: { + service_class: ::Gitlab::Llm::OpenAi::Completions::ExplainVulnerability, + prompt_class: ::Gitlab::Llm::OpenAi::Templates::ExplainVulnerability + }, summarize_comments: { service_class: ::Gitlab::Llm::OpenAi::Completions::SummarizeAllOpenNotes, prompt_class: ::Gitlab::Llm::OpenAi::Templates::SummarizeAllOpenNotes diff --git a/ee/lib/gitlab/llm/open_ai/response_service.rb b/ee/lib/gitlab/llm/open_ai/response_service.rb index 2d203a9b8e82eacf6e5ac7900752ddca04e4b6c9..b5df609989c49004bf18f64dc99a44ca0e1783dd 100644 --- a/ee/lib/gitlab/llm/open_ai/response_service.rb +++ b/ee/lib/gitlab/llm/open_ai/response_service.rb @@ -19,7 +19,7 @@ def execute(response_modifier = Gitlab::Llm::OpenAi::ResponseModifiers::Completi model_name: resource.class.name, # todo: do we need to sanitize/refine this response in any ways? response_body: response_modifier.execute(ai_response).to_s.strip, - errors: [ai_response&.dig(:error)].compact! + errors: [ai_response&.dig(:error)].compact } GraphqlTriggers.ai_completion_response(user.to_global_id, resource.to_global_id, data) diff --git a/ee/lib/gitlab/llm/open_ai/templates/explain_vulnerability.rb b/ee/lib/gitlab/llm/open_ai/templates/explain_vulnerability.rb new file mode 100644 index 0000000000000000000000000000000000000000..d04055bfad3d33b71e97153261585391fa41e540 --- /dev/null +++ b/ee/lib/gitlab/llm/open_ai/templates/explain_vulnerability.rb @@ -0,0 +1,86 @@ +# frozen_string_literal: true + +module Gitlab + module Llm + module OpenAi + module Templates + class ExplainVulnerability + include Gitlab::Utils::StrongMemoize + + def initialize(vulnerability) + @vulnerability = vulnerability + end + + def options + { + max_tokens: ::Llm::ExplainCodeService::MAX_RESPONSE_TOKENS + } + end + + def to_prompt + return prompt_without_file_or_code unless file.present? + return prompt_without_code unless source_code? + + default_prompt + end + + private + + attr_reader :vulnerability + + delegate :title, :description, :file, to: :vulnerability + delegate :source_code?, :vulnerable_code, to: :finding + + # rubocop: disable CodeReuse/ActiveRecord + def identifiers + finding.identifiers.pluck(:name).join(",") + end + # rubocop: enable CodeReuse/ActiveRecord + + def finding + vulnerability.finding + end + strong_memoize_attr :finding + + def filename + File.basename(file) + end + + def default_prompt + <<~PROMPT + You are a software vulnerability developer. + Explain the vulnerability "#{title} - #{description} (#{identifiers})". + The file "#{filename}" has this vulnerable code: + + ``` + #{vulnerable_code}" + ``` + + Provide a code example with syntax highlighting on how to exploit it. + Provide a code example with syntax highlighting on how to fix it. + PROMPT + end + + def prompt_without_code + <<~PROMPT + You are a software vulnerability developer. + Explain the vulnerability "#{title} - #{description} (#{identifiers})". + The vulnerable code is in the file "#{filename}". + Provide a code example with syntax highlighting on how to exploit it. + Provide a code example with syntax highlighting on how to fix it. + PROMPT + end + + def prompt_without_file_or_code + <<~PROMPT + You are a software vulnerability developer. + Explain the vulnerability "#{title} - #{description} (#{identifiers})". + Provide a code example with syntax highlighting on how to exploit it. + Provide a code example with syntax highlighting on how to fix it. + PROMPT + end + end + end + end + end +end diff --git a/ee/spec/lib/gitlab/llm/open_ai/completions/explain_vulnerability_spec.rb b/ee/spec/lib/gitlab/llm/open_ai/completions/explain_vulnerability_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..27ddc8f6f41329be9710fdcf948a2f1c2896c5a3 --- /dev/null +++ b/ee/spec/lib/gitlab/llm/open_ai/completions/explain_vulnerability_spec.rb @@ -0,0 +1,137 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Llm::OpenAi::Completions::ExplainVulnerability, feature_category: :vulnerability_management do + let(:prompt_class) { Gitlab::Llm::OpenAi::Templates::ExplainVulnerability } + let_it_be(:user) { create(:user) } + let_it_be(:project) do + create(:project, :custom_repo, files: { + 'main.c' => "#include <stdio.h>\n\nint main() { printf(\"hello, world!\"); }" + }) + end + + let_it_be(:vulnerability) { create(:vulnerability, :with_finding, project: project) } + + subject(:explain) { described_class.new(prompt_class) } + + before do + allow(GraphqlTriggers).to receive(:ai_completion_response) + vulnerability.finding.location['file'] = 'main.c' + vulnerability.finding.location['start_line'] = 1 + end + + describe '#execute' do + context 'when the chat client returns an unsuccessful response' do + before do + allow_next_instance_of(Gitlab::Llm::OpenAi::Client) do |client| + allow(client).to receive(:chat).and_return( + { 'error' => 'Ooops...' }.to_json + ) + end + end + + it 'publishes the error to the graphql subscription' do + explain.execute(user, vulnerability, {}) + + expect(GraphqlTriggers).to have_received(:ai_completion_response) + .with(user.to_global_id, vulnerability.to_global_id, hash_including({ + id: anything, + model_name: vulnerability.class.name, + response_body: '', + errors: ['Ooops...'] + })) + end + end + + context 'when the chat client returns a successful response' do + let(:example_answer) do + <<-AI.strip + As an AI language model, I cannot access or analyze specific files or code. + However, based on the limited information provided, "Test vulnerability 1" + could refer to a security flaw in a particular software or application. + It is likely that the vulnerability could be exploited by an attacker to gain + unauthorized access to sensitive data or execute malicious code. + + To fix this vulnerability, the code in question needs to be analyzed and modified + to remove any security flaws. The "Test remediations" section of the README.md + file should provide step-by-step instructions on how to fix the vulnerability, + which may involve updating the software or application, applying security patches, + or implementing additional security measures. It is important to follow these + instructions carefully to ensure that the vulnerability is fully + addressed and the system remains secure. + AI + end + + let(:example_response) do + { + 'id' => 'chatcmpl-74uDpPnYHVPwLg0RIM6recPqgZKm5', + 'object' => 'chat.completion', + 'created' => 1681403785, + 'model' => 'gpt-3.5-turbo-0301', + 'usage' => { + 'prompt_tokens' => 59, + 'completion_tokens' => 155, + 'total_tokens' => 214 + }, + 'choices' => [ + { + 'message' => { + 'role' => 'assistant', + 'content' => example_answer + }, + 'finish_reason' => 'stop', + 'index' => 0 + } + ] + } + end + + before do + allow_next_instance_of(Gitlab::Llm::OpenAi::Client) do |client| + allow(client).to receive(:chat).and_return(example_response.to_json) + end + end + + it 'publishes the content field from the AI response' do + explain.execute(user, vulnerability, {}) + + expect(GraphqlTriggers).to have_received(:ai_completion_response) + .with(user.to_global_id, vulnerability.to_global_id, hash_including({ + id: anything, + model_name: vulnerability.class.name, + response_body: example_answer, + errors: [] + })) + end + + context 'when an unexpected error is raised' do + let(:error) { StandardError.new("Ooops...") } + + before do + allow_next_instance_of(Gitlab::Llm::OpenAi::Client) do |client| + allow(client).to receive(:chat).and_raise(error) + end + allow(Gitlab::ErrorTracking).to receive(:track_exception) + end + + it 'records the error' do + explain.execute(user, vulnerability, {}) + expect(Gitlab::ErrorTracking).to have_received(:track_exception).with(error) + end + + it 'publishes a generic error to the graphql subscription' do + explain.execute(user, vulnerability, {}) + + expect(GraphqlTriggers).to have_received(:ai_completion_response) + .with(user.to_global_id, vulnerability.to_global_id, hash_including({ + id: anything, + model_name: vulnerability.class.name, + response_body: '', + errors: ['An unexpected error has occurred.'] + })) + end + end + end + end +end diff --git a/ee/spec/lib/gitlab/llm/open_ai/templates/explain_vulnerability_spec.rb b/ee/spec/lib/gitlab/llm/open_ai/templates/explain_vulnerability_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..2fa8edceec7eeffe524f77f7313ef9334b7e772a --- /dev/null +++ b/ee/spec/lib/gitlab/llm/open_ai/templates/explain_vulnerability_spec.rb @@ -0,0 +1,104 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Llm::OpenAi::Templates::ExplainVulnerability, feature_category: :vulnerability_management do + let_it_be(:source_code) do + <<~SOURCE + #include <stdio.h> + + int main(int argc, char *argv[]) + { + char buf[8]; + memcpy(&buf, "123456789"); + printf("hello, world!"); + } + SOURCE + end + + let_it_be(:project) do + create(:project, :custom_repo, files: { + 'src/main.c' => source_code + }) + end + + let_it_be(:vulnerability) do + create(:vulnerability, :with_finding, project: project) + end + + before do + vulnerability.finding.project = project + vulnerability.finding.clear_memoization(:source_code) + vulnerability.finding.clear_memoization(:vulnerable_code) + end + + subject { described_class.new(vulnerability) } + + describe '#to_prompt' do + let(:identifiers) { vulnerability.finding.identifiers.pluck(:name).join(",") } + + context 'when a file is provided' do + context 'when the file exists' do + before do + vulnerability.finding.location['file'] = 'src/main.c' + vulnerability.finding.location['start_line'] = 5 + vulnerability.finding.location['end_line'] = 6 + end + + it 'includes the title' do + expect(subject.to_prompt).to include(vulnerability.title) + end + + it 'includes the description' do + expect(subject.to_prompt).to include(vulnerability.description) + end + + it 'includes the identifiers' do + expect(subject.to_prompt).to include(identifiers) + end + + it 'includes the file name' do + expect(subject.to_prompt).to include('"main.c"') + end + + it 'includes the vulnerable code' do + vulnerable_code = source_code.lines[4..5].join + expect(subject.to_prompt).to include(vulnerable_code) + end + end + + context 'when the file does not exist' do + before do + vulnerability.finding.location['file'] = 'missing.c' + end + + it 'customizes the prompt' do + expect(subject.to_prompt).to eq(<<~PROMPT) + You are a software vulnerability developer. + Explain the vulnerability "#{vulnerability.title} - #{vulnerability.description} (#{identifiers})". + The vulnerable code is in the file "#{vulnerability.file}". + Provide a code example with syntax highlighting on how to exploit it. + Provide a code example with syntax highlighting on how to fix it. + PROMPT + end + end + end + + context 'when a file is not provided' do + before do + vulnerability.finding.location.delete('file') + end + + it 'customizes the prompt' do + expected = <<~PROMPT + You are a software vulnerability developer. + Explain the vulnerability "#{vulnerability.title} - #{vulnerability.description} (#{identifiers})". + Provide a code example with syntax highlighting on how to exploit it. + Provide a code example with syntax highlighting on how to fix it. + PROMPT + + expect(subject.to_prompt).to eq(expected) + end + end + end +end diff --git a/ee/spec/models/ee/vulnerability_spec.rb b/ee/spec/models/ee/vulnerability_spec.rb index f08d3730a5965765e4da8fca2d3d731a8be2c923..01c895955c60ed935a6ade2a692b0e6ce2c45b00 100644 --- a/ee/spec/models/ee/vulnerability_spec.rb +++ b/ee/spec/models/ee/vulnerability_spec.rb @@ -929,4 +929,56 @@ end end end + + describe '#send_to_ai?' do + context 'when the feature flag is disabled' do + let_it_be(:vulnerability) { create(:vulnerability) } + + before do + stub_feature_flags(explain_vulnerability: false) + end + + it 'returns false' do + expect(vulnerability).not_to be_send_to_ai + end + end + + context 'when the feature flag is enabled' do + before do + stub_feature_flags(explain_vulnerability: vulnerability.project) + end + + context 'when the vulnerabilty is from SAST and includes the required information' do + let_it_be(:vulnerability) { create(:vulnerability, :with_finding) } + + it { expect(vulnerability).to be_send_to_ai } + end + + context 'when the vulnerability does not include a file' do + let_it_be(:vulnerability) { create(:vulnerability, :with_finding) } + + before do + vulnerability.finding.location.delete('file') + end + + it { expect(vulnerability).not_to be_send_to_ai } + end + + context 'when the vulnerability does not include a start line' do + let_it_be(:vulnerability) { create(:vulnerability, :with_finding) } + + before do + vulnerability.finding.location.delete('start_line') + end + + it { expect(vulnerability).not_to be_send_to_ai } + end + + context 'when the vulnerability does not include a finding' do + let_it_be(:vulnerability) { create(:vulnerability) } + + it { expect(vulnerability).not_to be_send_to_ai } + end + end + end end diff --git a/ee/spec/models/vulnerabilities/finding_spec.rb b/ee/spec/models/vulnerabilities/finding_spec.rb index 28fa52b1a816a803f96f7b9eee8f42d8d58ed37d..31f0419cc31a41f310ed24f01ab7bbce050816fd 100644 --- a/ee/spec/models/vulnerabilities/finding_spec.rb +++ b/ee/spec/models/vulnerabilities/finding_spec.rb @@ -1133,4 +1133,44 @@ it { is_expected.to contain_exactly(finding_2) } end + + describe "#vulnerable_code" do + let_it_be(:source_code) do + <<~SOURCE + #include <stdio.h> + + int main(int argc, char *argv[]) + { + char buf[8]; + memcpy(&buf, "123456789"); + printf("hello, world!"); + } + SOURCE + end + + let_it_be(:project) do + create(:project, :custom_repo, files: { + 'src/main.c' => source_code + }) + end + + let_it_be(:finding) do + create(:vulnerabilities_finding).tap do |finding| + finding.project = project + finding.location['file'] = 'src/main.c' + finding.location['start_line'] = 5 + finding.location['end_line'] = 6 + end + end + + subject { finding.vulnerable_code } + + it 'returns the vulnerables lines of code' do + vulnerable_lines = <<-LINES + char buf[8]; + memcpy(&buf, "123456789"); + LINES + expect(subject).to eq(vulnerable_lines) + end + end end diff --git a/ee/spec/policies/vulnerability_policy_spec.rb b/ee/spec/policies/vulnerability_policy_spec.rb index 82f8da8c8d12801e5e63fff0db01ea4e7d3751bb..25899da937800a4f187852e41d96ec791b62346a 100644 --- a/ee/spec/policies/vulnerability_policy_spec.rb +++ b/ee/spec/policies/vulnerability_policy_spec.rb @@ -2,7 +2,7 @@ require 'spec_helper' -RSpec.describe VulnerabilityPolicy do +RSpec.describe VulnerabilityPolicy, feature_category: :vulnerability_management do describe 'read_security_resource' do let(:project) { create(:project) } let(:user) { create(:user) } @@ -17,6 +17,7 @@ context "when the current user is not a project member" do it { is_expected.to be_disallowed(:read_security_resource) } + it { is_expected.to be_disallowed(:read_vulnerability) } it { is_expected.to be_disallowed(:create_note) } end @@ -26,6 +27,7 @@ end it { is_expected.to be_allowed(:read_security_resource) } + it { is_expected.to be_allowed(:read_vulnerability) } it { is_expected.to be_allowed(:create_note) } end @@ -46,6 +48,7 @@ end it { is_expected.to be_disallowed(:read_security_resource) } + it { is_expected.to be_disallowed(:read_vulnerability) } it { is_expected.to be_disallowed(:create_note) } end end diff --git a/ee/spec/services/llm/execute_method_service_spec.rb b/ee/spec/services/llm/execute_method_service_spec.rb index f3789385e33a08f293f94a1415cf8350c040cd70..6b6cc3da0807b6a30e509f5b83f1544521310c13 100644 --- a/ee/spec/services/llm/execute_method_service_spec.rb +++ b/ee/spec/services/llm/execute_method_service_spec.rb @@ -19,6 +19,7 @@ where(:method, :resource, :service_class) do :summarize_comments | issue | Llm::GenerateSummaryService :explain_code | build_stubbed(:project) | Llm::ExplainCodeService + :explain_vulnerability | build_stubbed(:vulnerability, :with_findings) | Llm::ExplainVulnerabilityService end with_them do diff --git a/ee/spec/services/llm/explain_vulnerability_service_spec.rb b/ee/spec/services/llm/explain_vulnerability_service_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..ec0f62290c98c675229dc09627d801a328ab29e1 --- /dev/null +++ b/ee/spec/services/llm/explain_vulnerability_service_spec.rb @@ -0,0 +1,65 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Llm::ExplainVulnerabilityService, feature_category: :vulnerability_management do + let_it_be(:user) { create(:user) } + let_it_be(:project) { create(:project, :private) } + let_it_be(:vulnerability) { create(:vulnerability, :with_finding, project: project) } + let_it_be(:options) { {} } + + subject { described_class.new(user, vulnerability, options) } + + before do + stub_feature_flags(openai_experimentation: true) + end + + describe '#execute' do + before do + stub_feature_flags(explain_vulnerability: project) + allow(Llm::CompletionWorker).to receive(:perform_async) + end + + context 'when the user is permitted to view the vulnerability' do + before do + project.add_maintainer(user) + end + + it 'schedules a job' do + expect(subject.execute).to be_success + + expect(Llm::CompletionWorker).to have_received(:perform_async).with( + user.id, + vulnerability.id, + 'Vulnerability', + :explain_vulnerability, + options + ) + end + end + + context 'when the user is not permitted to view the vulnerability' do + before do + allow(project).to receive(:member?).with(user).and_return(false) + end + + it 'returns an error' do + expect(subject.execute).to be_error + + expect(Llm::CompletionWorker).not_to have_received(:perform_async) + end + end + + context 'when feature flag is disabled' do + before do + stub_feature_flags(explain_vulnerability: false) + end + + it 'returns an error' do + expect(subject.execute).to be_error + + expect(Llm::CompletionWorker).not_to have_received(:perform_async) + end + end + end +end