Add work item mutation arguments for custom fields

Allows updating custom field values of a work item

app/graphql/mutations/concerns/mutations/work_items/widgetable.rb

@@ -16,9 +16,11 @@ def extract_widget_params!(work_item_type, attributes, resource_parent)
             "Following widget keys are not supported by #{work_item_type.name} type: #{not_supported_keys}"
         end
 
-        # Cannot use prepare to use `.to_h` on each input due to
-        # https://gitlab.com/gitlab-org/gitlab/-/merge_requests/87472#note_945199865
-        widget_params.transform_values(&:to_h)
+        # TODO: Refactor to use `#prepare` on the input types
+        # https://gitlab.com/gitlab-org/gitlab/-/issues/519801
+        widget_params.transform_values do |input|
+          input.is_a?(Array) ? input.map(&:to_h) : input.to_h
+        end
       end
     end
   end

doc/api/graphql/reference/_index.md

@@ -12233,6 +12233,7 @@ Input type: `WorkItemCreateInput`
 | <a id="mutationworkitemcreateconfidential"></a>`confidential` | [`Boolean`](#boolean) | Sets the work item confidentiality. |
 | <a id="mutationworkitemcreatecreatedat"></a>`createdAt` | [`Time`](#time) | Timestamp when the work item was created. Available only for admins and project owners. |
 | <a id="mutationworkitemcreatecrmcontactswidget"></a>`crmContactsWidget` | [`WorkItemWidgetCrmContactsCreateInput`](#workitemwidgetcrmcontactscreateinput) | Input for CRM contacts widget. |
+| <a id="mutationworkitemcreatecustomfieldswidget"></a>`customFieldsWidget` {{< icon name="warning-solid" >}} | [`[WorkItemWidgetCustomFieldValueInputType!]`](#workitemwidgetcustomfieldvalueinputtype) | **Deprecated:** **Status**: Experiment. Introduced in GitLab 17.10. |
 | <a id="mutationworkitemcreatedescription"></a>`description` {{< icon name="warning-solid" >}} | [`String`](#string) | **Deprecated:** use description widget instead. Deprecated in GitLab 16.9. |
 | <a id="mutationworkitemcreatedescriptionwidget"></a>`descriptionWidget` | [`WorkItemWidgetDescriptionInput`](#workitemwidgetdescriptioninput) | Input for description widget. |
 | <a id="mutationworkitemcreatediscussionstoresolve"></a>`discussionsToResolve` | [`WorkItemResolveDiscussionsInput`](#workitemresolvediscussionsinput) | Information required to resolve discussions in a noteable, when the work item is created. |
@@ -12437,6 +12438,7 @@ Input type: `WorkItemUpdateInput`
 | <a id="mutationworkitemupdateconfidential"></a>`confidential` | [`Boolean`](#boolean) | Sets the work item confidentiality. |
 | <a id="mutationworkitemupdatecrmcontactswidget"></a>`crmContactsWidget` | [`WorkItemWidgetCrmContactsUpdateInput`](#workitemwidgetcrmcontactsupdateinput) | Input for CRM contacts widget. |
 | <a id="mutationworkitemupdatecurrentusertodoswidget"></a>`currentUserTodosWidget` | [`WorkItemWidgetCurrentUserTodosInput`](#workitemwidgetcurrentusertodosinput) | Input for to-dos widget. |
+| <a id="mutationworkitemupdatecustomfieldswidget"></a>`customFieldsWidget` {{< icon name="warning-solid" >}} | [`[WorkItemWidgetCustomFieldValueInputType!]`](#workitemwidgetcustomfieldvalueinputtype) | **Deprecated:** **Status**: Experiment. Introduced in GitLab 17.10. |
 | <a id="mutationworkitemupdatedescriptionwidget"></a>`descriptionWidget` | [`WorkItemWidgetDescriptionInput`](#workitemwidgetdescriptioninput) | Input for description widget. |
 | <a id="mutationworkitemupdatehealthstatuswidget"></a>`healthStatusWidget` | [`WorkItemWidgetHealthStatusInput`](#workitemwidgethealthstatusinput) | Input for health status widget. |
 | <a id="mutationworkitemupdatehierarchywidget"></a>`hierarchyWidget` | [`WorkItemWidgetHierarchyUpdateInput`](#workitemwidgethierarchyupdateinput) | Input for hierarchy widget. |
@@ -47197,6 +47199,17 @@ Attributes for value stream stage.
 | <a id="workitemwidgetcurrentusertodosinputaction"></a>`action` | [`WorkItemTodoUpdateAction!`](#workitemtodoupdateaction) | Action for the update. |
 | <a id="workitemwidgetcurrentusertodosinputtodoid"></a>`todoId` | [`TodoID`](#todoid) | Global ID of the to-do. If not present, all to-dos of the work item will be updated. |
+### `WorkItemWidgetCustomFieldValueInputType`
+
+#### Arguments
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="workitemwidgetcustomfieldvalueinputtypecustomfieldid"></a>`customFieldId` | [`IssuablesCustomFieldID!`](#issuablescustomfieldid) | Global ID of the custom field. |
+| <a id="workitemwidgetcustomfieldvalueinputtypenumbervalue"></a>`numberValue` | [`Float`](#float) | Value for custom fields with number type. |
+| <a id="workitemwidgetcustomfieldvalueinputtypeselectedoptionids"></a>`selectedOptionIds` | [`[IssuablesCustomFieldSelectOptionID!]`](#issuablescustomfieldselectoptionid) | Global IDs of the selected options for custom fields with select type. |
+| <a id="workitemwidgetcustomfieldvalueinputtypetextvalue"></a>`textValue` | [`String`](#string) | Value for custom fields with text type. |
+
 ### `WorkItemWidgetDescriptionInput`
 #### Arguments

ee/app/graphql/ee/mutations/work_items/create.rb

@@ -27,6 +27,11 @@ module Create
             required: false,
             description: 'Input for color widget.'
 
+          argument :custom_fields_widget, [::Types::WorkItems::Widgets::CustomFieldValueInputType],
+            required: false,
+            description: 'Input for custom fields widget.',
+            experiment: { milestone: '17.10' }
+
           argument :vulnerability_id, ::Types::GlobalIDType[::Vulnerability],
             required: false,
             description: 'Input for linking an existing vulnerability to created work item.',

ee/app/graphql/ee/mutations/work_items/update.rb

@@ -30,6 +30,11 @@ module Update
           argument :color_widget, ::Types::WorkItems::Widgets::ColorInputType,
             required: false,
             description: 'Input for color widget.'
+
+          argument :custom_fields_widget, [::Types::WorkItems::Widgets::CustomFieldValueInputType],
+            required: false,
+            description: 'Input for custom fields widget.',
+            experiment: { milestone: '17.10' }
         end
       end
     end

ee/app/graphql/types/work_items/widgets/custom_field_value_input_type.rb

+# frozen_string_literal: true
+
+module Types
+  module WorkItems
+    module Widgets
+      class CustomFieldValueInputType < BaseInputObject
+        graphql_name 'WorkItemWidgetCustomFieldValueInputType'
+
+        argument :custom_field_id, ::Types::GlobalIDType[::Issuables::CustomField],
+          required: true,
+          description: copy_field_description(Types::Issuables::CustomFieldType, :id),
+          prepare: ->(id, _ctx) { id&.model_id }
+
+        argument :selected_option_ids, [::Types::GlobalIDType[::Issuables::CustomFieldSelectOption]],
+          required: false,
+          description: 'Global IDs of the selected options for custom fields with select type.',
+          prepare: ->(ids, _ctx) { ids.map(&:model_id) }
+
+        argument :number_value, GraphQL::Types::Float,
+          required: false,
+          description: 'Value for custom fields with number type.'
+
+        argument :text_value, GraphQL::Types::String,
+          required: false,
+          description: 'Value for custom fields with text type.'
+
+        validates mutually_exclusive: [:selected_option_ids, :number_value, :text_value]
+      end
+    end
+  end
+end

ee/app/models/concerns/work_items/scalar_custom_field_value.rb

+# frozen_string_literal: true
+
+module WorkItems
+  module ScalarCustomFieldValue
+    extend ActiveSupport::Concern
+
+    class_methods do
+      def update_work_item_field!(work_item, field, value)
+        return where(work_item: work_item, custom_field: field).delete_all if value.blank?
+
+        field_value = find_or_initialize_by(work_item: work_item, custom_field: field)
+        field_value.update!(value: value)
+      rescue ActiveRecord::RecordInvalid => invalid
+        retry if invalid.record&.errors&.of_kind?(:custom_field, :taken)
+        raise
+      end
+    end
+  end
+end

ee/app/models/work_items/number_field_value.rb

@@ -3,6 +3,7 @@
 module WorkItems
   class NumberFieldValue < ApplicationRecord
     include CustomFieldValue
+    include ScalarCustomFieldValue
 
     validates :custom_field, uniqueness: { scope: [:work_item_id] }
     validates :value, presence: true, numericality: true

ee/app/models/work_items/select_field_value.rb

@@ -7,5 +7,58 @@ class SelectFieldValue < ApplicationRecord
     belongs_to :custom_field_select_option, class_name: 'Issuables::CustomFieldSelectOption'
 
     validates :custom_field_select_option, presence: true, uniqueness: { scope: [:work_item_id, :custom_field_id] }
+
+    class << self
+      def update_work_item_field!(work_item, field, selected_option_ids)
+        return where(work_item: work_item, custom_field: field).delete_all if selected_option_ids.blank?
+
+        selected_option_ids = selected_option_ids.uniq.map(&:to_i)
+
+        if field.field_type_single_select? && selected_option_ids.size > 1
+          raise ArgumentError, 'A custom field of type single select may only have a single selected option'
+        end
+
+        existing_option_ids = where(work_item: work_item, custom_field: field).pluck(:custom_field_select_option_id) # rubocop:disable Database/AvoidUsingPluckWithoutLimit -- the number of options for a select field is limited
+
+        transaction do
+          delete_ids(work_item: work_item, field: field, ids: existing_option_ids - selected_option_ids)
+          insert_ids(work_item: work_item, field: field, ids: selected_option_ids - existing_option_ids)
+        end
+      end
+
+      private
+
+      def delete_ids(work_item:, field:, ids:)
+        return if ids.empty?
+
+        where(
+          work_item: work_item,
+          custom_field: field,
+          custom_field_select_option_id: ids
+        ).delete_all
+      end
+
+      def insert_ids(work_item:, field:, ids:)
+        return if ids.empty?
+
+        available_option_ids = field.select_options.id_in(ids).pluck_primary_key
+
+        if available_option_ids.size != ids.size
+          raise ArgumentError,
+            "Invalid custom field select option IDs: #{(ids - available_option_ids).join(',')}"
+        end
+
+        attributes_to_insert = ids.map do |option_id|
+          {
+            namespace_id: work_item.namespace_id,
+            work_item_id: work_item.id,
+            custom_field_id: field.id,
+            custom_field_select_option_id: option_id
+          }
+        end
+
+        insert_all(attributes_to_insert)
+      end
+    end
   end
 end

ee/app/models/work_items/text_field_value.rb

@@ -3,6 +3,7 @@
 module WorkItems
   class TextFieldValue < ApplicationRecord
     include CustomFieldValue
+    include ScalarCustomFieldValue
 
     MAX_LENGTH = 1024
 

ee/app/services/work_items/callbacks/custom_fields.rb

+# frozen_string_literal: true
+
+module WorkItems
+  module Callbacks
+    class CustomFields < Base
+      # `params` for this widget callback is in the format:
+      # [
+      #   { custom_field_id: 1, text_value: 'text' },
+      #   { custom_field_id: 2, number_value: 100 },
+      #   { custom_field_id: 3, selected_option_ids: [1, 2, 3] }
+      # ]
+      # Only values for the provided custom_field_ids are mutated. Omitted ones are left as-is.
+      def after_save
+        return unless Feature.enabled?(:custom_fields_feature, work_item.namespace.root_ancestor)
+        return unless has_permission?(:set_work_item_metadata)
+
+        custom_fields = ::Issuables::CustomFieldsFinder.active_fields_for_work_item(work_item)
+                          .id_in(params.pluck(:custom_field_id)) # rubocop:disable CodeReuse/ActiveRecord, Database/AvoidUsingPluckWithoutLimit -- params is an Array
+                          .index_by(&:id)
+
+        params.each do |field_params|
+          custom_field = custom_fields[field_params[:custom_field_id].to_i]
+
+          raise_error "Invalid custom field ID: #{field_params[:custom_field_id]}" if custom_field.nil?
+
+          update_work_item_field_value(custom_field, field_params)
+        end
+      end
+
+      private
+
+      def update_work_item_field_value(custom_field, field_params)
+        if custom_field.field_type_text?
+          WorkItems::TextFieldValue
+            .update_work_item_field!(work_item, custom_field, field_params[:text_value])
+        elsif custom_field.field_type_number?
+          WorkItems::NumberFieldValue
+            .update_work_item_field!(work_item, custom_field, field_params[:number_value])
+        elsif custom_field.field_type_select?
+          WorkItems::SelectFieldValue
+            .update_work_item_field!(work_item, custom_field, field_params[:selected_option_ids])
+        else
+          raise_error "Unsupported field type: #{custom_field.field_type}"
+        end
+      rescue ActiveRecord::RecordInvalid, ArgumentError => e
+        raise_error e.message
+      end
+    end
+  end
+end

ee/spec/graphql/types/work_items/widgets/custom_field_value_input_type_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ::Types::WorkItems::Widgets::CustomFieldValueInputType, feature_category: :team_planning do
+  it { expect(described_class.graphql_name).to eq('WorkItemWidgetCustomFieldValueInputType') }
+
+  it 'has correct arguments' do
+    expect(described_class.arguments.keys).to contain_exactly(
+      'customFieldId', 'selectedOptionIds', 'numberValue', 'textValue'
+    )
+  end
+end

ee/spec/models/work_items/number_field_value_spec.rb

@@ -11,4 +11,69 @@
     it { is_expected.to validate_presence_of(:value) }
     it { is_expected.to validate_numericality_of(:value) }
   end
+
+  describe '.update_work_item_field!' do
+    let_it_be(:group) { create(:group) }
+    let_it_be(:project) { create(:project, group: group) }
+
+    let_it_be(:work_item) { create(:work_item, project: project) }
+    let_it_be(:custom_field) do
+      create(:custom_field, namespace: group, field_type: 'number', work_item_types: [work_item.work_item_type])
+    end
+
+    context 'when there is no existing record' do
+      it 'inserts a new record' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, 100)
+        end.to change { described_class.count }.by(1)
+
+        expect(described_class.last).to have_attributes({
+          work_item_id: work_item.id,
+          custom_field_id: custom_field.id,
+          value: 100
+        })
+      end
+
+      it 'retries in case of a race condition' do
+        expect_next_instance_of(described_class) do |field_value|
+          field_value.errors.add(:custom_field, :taken)
+          expect(field_value).to receive(:update!).and_raise(ActiveRecord::RecordInvalid.new(field_value))
+        end
+
+        expect_next_instance_of(described_class) do |field_value|
+          expect(field_value).to receive(:update!).and_call_original
+        end
+
+        described_class.update_work_item_field!(work_item, custom_field, 100)
+      end
+
+      context 'when there is a validation error' do
+        it 'raises an error' do
+          expect do
+            described_class.update_work_item_field!(work_item, custom_field, 'some string')
+          end.to raise_error(ActiveRecord::RecordInvalid, 'Validation failed: Value is not a number')
+        end
+      end
+    end
+
+    context 'when there is an existing record' do
+      let!(:existing_field_value) do
+        create(:work_item_number_field_value, work_item: work_item, custom_field: custom_field, value: 50)
+      end
+
+      it 'updates the existing record' do
+        described_class.update_work_item_field!(work_item, custom_field, 100)
+
+        expect(existing_field_value.reload.value).to eq(100)
+      end
+
+      it 'deletes the record when value is set to nil' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, nil)
+        end.to change { described_class.count }.by(-1)
+
+        expect { existing_field_value.reload }.to raise_error(ActiveRecord::RecordNotFound)
+      end
+    end
+  end
 end

ee/spec/models/work_items/select_field_value_spec.rb

@@ -21,4 +21,102 @@
       is_expected.to validate_uniqueness_of(:custom_field_select_option).scoped_to([:work_item_id, :custom_field_id])
     end
   end
+
+  describe '.update_work_item_field!' do
+    let_it_be(:group) { create(:group) }
+    let_it_be(:project) { create(:project, group: group) }
+
+    let_it_be(:work_item) { create(:work_item, project: project) }
+    let_it_be(:custom_field) do
+      create(:custom_field, namespace: group, field_type: 'multi_select', work_item_types: [work_item.work_item_type])
+    end
+
+    let_it_be(:multi_select_option_1) { create(:custom_field_select_option, custom_field: custom_field) }
+    let_it_be(:multi_select_option_2) { create(:custom_field_select_option, custom_field: custom_field) }
+    let_it_be(:multi_select_option_3) { create(:custom_field_select_option, custom_field: custom_field) }
+
+    context 'when there are no existing records' do
+      it 'inserts a new record for each selected option' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, [
+            multi_select_option_1, multi_select_option_3
+          ].map(&:id))
+        end.to change { described_class.count }.by(2)
+
+        expect(described_class.last(2)).to contain_exactly(
+          have_attributes({
+            work_item_id: work_item.id,
+            custom_field_id: custom_field.id,
+            custom_field_select_option_id: multi_select_option_1.id
+          }),
+          have_attributes({
+            work_item_id: work_item.id,
+            custom_field_id: custom_field.id,
+            custom_field_select_option_id: multi_select_option_3.id
+          })
+        )
+      end
+
+      it 'raises an argument error when passing a select option of an unrelated custom field' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, [
+            create(:custom_field_select_option).id
+          ])
+        end.to raise_error(ArgumentError, /Invalid custom field select option IDs/)
+      end
+    end
+
+    context 'when there are existing records' do
+      before do
+        create(:work_item_select_field_value, work_item: work_item, custom_field: custom_field,
+          custom_field_select_option: multi_select_option_2)
+      end
+
+      it 'inserts and deletes records to match the new selected options' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, [
+            multi_select_option_1, multi_select_option_3
+          ].map(&:id))
+        end.to change { described_class.count }.by(1)
+
+        expect(described_class.last(2)).to contain_exactly(
+          have_attributes({
+            work_item_id: work_item.id,
+            custom_field_id: custom_field.id,
+            custom_field_select_option_id: multi_select_option_1.id
+          }),
+          have_attributes({
+            work_item_id: work_item.id,
+            custom_field_id: custom_field.id,
+            custom_field_select_option_id: multi_select_option_3.id
+          })
+        )
+      end
+
+      it 'deletes existing records when set to nil' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, nil)
+        end.to change { described_class.count }.by(-1)
+      end
+    end
+
+    context 'when custom field is a single select' do
+      let_it_be(:custom_field) do
+        create(:custom_field, namespace: group, field_type: 'single_select',
+          work_item_types: [work_item.work_item_type])
+      end
+
+      let_it_be(:select_option_1) { create(:custom_field_select_option, custom_field: custom_field) }
+      let_it_be(:select_option_2) { create(:custom_field_select_option, custom_field: custom_field) }
+
+      it 'raises an argument error when passing multiple options' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, [
+            select_option_1,
+            select_option_2
+          ].map(&:id))
+        end.to raise_error(ArgumentError, 'A custom field of type single select may only have a single selected option')
+      end
+    end
+  end
 end

ee/spec/models/work_items/text_field_value_spec.rb

@@ -11,4 +11,70 @@
     it { is_expected.to validate_presence_of(:value) }
     it { is_expected.to validate_length_of(:value).is_at_most(described_class::MAX_LENGTH) }
   end
+
+  describe '.update_work_item_field!' do
+    let_it_be(:group) { create(:group) }
+    let_it_be(:project) { create(:project, group: group) }
+
+    let_it_be(:work_item) { create(:work_item, project: project) }
+    let_it_be(:custom_field) do
+      create(:custom_field, namespace: group, field_type: 'text', work_item_types: [work_item.work_item_type])
+    end
+
+    context 'when there is no existing record' do
+      it 'inserts a new record' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, 'some text')
+        end.to change { described_class.count }.by(1)
+
+        expect(described_class.last).to have_attributes({
+          work_item_id: work_item.id,
+          custom_field_id: custom_field.id,
+          value: 'some text'
+        })
+      end
+
+      it 'retries in case of a race condition' do
+        expect_next_instance_of(described_class) do |field_value|
+          field_value.errors.add(:custom_field, :taken)
+          expect(field_value).to receive(:update!).and_raise(ActiveRecord::RecordInvalid.new(field_value))
+        end
+
+        expect_next_instance_of(described_class) do |field_value|
+          expect(field_value).to receive(:update!).and_call_original
+        end
+
+        described_class.update_work_item_field!(work_item, custom_field, 'some text')
+      end
+
+      context 'when there is a validation error' do
+        it 'raises an error' do
+          expect do
+            described_class.update_work_item_field!(work_item, custom_field, 'a' * (described_class::MAX_LENGTH + 1))
+          end.to raise_error(ActiveRecord::RecordInvalid,
+            'Validation failed: Value is too long (maximum is 1024 characters)')
+        end
+      end
+    end
+
+    context 'when there is an existing record' do
+      let!(:existing_field_value) do
+        create(:work_item_text_field_value, work_item: work_item, custom_field: custom_field, value: 'existing value')
+      end
+
+      it 'updates the existing record' do
+        described_class.update_work_item_field!(work_item, custom_field, 'some text')
+
+        expect(existing_field_value.reload.value).to eq('some text')
+      end
+
+      it 'deletes the record when value is set to nil' do
+        expect do
+          described_class.update_work_item_field!(work_item, custom_field, nil)
+        end.to change { described_class.count }.by(-1)
+
+        expect { existing_field_value.reload }.to raise_error(ActiveRecord::RecordNotFound)
+      end
+    end
+  end
 end

ee/spec/requests/api/graphql/mutations/work_items/create_spec.rb

@@ -774,4 +774,54 @@
       end
     end
   end
+
+  context 'with custom fields widget input' do
+    include_context 'with group configured with custom fields'
+
+    let_it_be(:project) { create(:project, group: group, developers: [developer]) }
+
+    let(:mutation) { graphql_mutation(:workItemCreate, input, 'workItem { id }') }
+
+    let(:input) do
+      {
+        'namespacePath' => project.full_path,
+        'workItemTypeId' => issue_type.to_gid.to_s,
+        'title' => 'New work item',
+        'customFieldsWidget' => [
+          { 'customFieldId' => global_id_of(number_field), 'numberValue' => 100 },
+          { 'customFieldId' => global_id_of(multi_select_field), 'selectedOptionIds' => [
+            global_id_of(multi_select_option_1), global_id_of(multi_select_option_3)
+          ] }
+        ]
+      }
+    end
+
+    before do
+      stub_licensed_features(custom_fields: true)
+    end
+
+    it 'creates the work item with custom field values' do
+      post_graphql_mutation(mutation, current_user: developer)
+
+      expect(response).to have_gitlab_http_status(:success)
+
+      work_item_id = GlobalID.parse(mutation_response['workItem']['id']).model_id.to_i
+
+      expect(WorkItems::NumberFieldValue.last).to have_attributes(
+        work_item_id: work_item_id, custom_field_id: number_field.id, value: 100
+      )
+      expect(WorkItems::SelectFieldValue.last(2)).to contain_exactly(
+        have_attributes({
+          work_item_id: work_item_id,
+          custom_field_id: multi_select_field.id,
+          custom_field_select_option_id: multi_select_option_1.id
+        }),
+        have_attributes({
+          work_item_id: work_item_id,
+          custom_field_id: multi_select_field.id,
+          custom_field_select_option_id: multi_select_option_3.id
+        })
+      )
+    end
+  end
 end

ee/spec/requests/api/graphql/mutations/work_items/update_spec.rb

@@ -1208,6 +1208,46 @@ def work_item_status
     end
   end
 
+  context 'with custom fields widget input' do
+    include_context 'with group configured with custom fields'
+
+    let_it_be(:project) { create(:project, group: group, reporters: [reporter]) }
+    let_it_be(:work_item) { create(:work_item, work_item_type: issue_type, project: project) }
+
+    let(:fields) { 'workItem { id }' }
+    let(:input) do
+      {
+        'customFieldsWidget' => [
+          { 'customFieldId' => global_id_of(text_field), 'textValue' => 'some text' },
+          { 'customFieldId' => global_id_of(select_field), 'selectedOptionIds' => [
+            global_id_of(select_option_1)
+          ] }
+        ]
+      }
+    end
+
+    before do
+      stub_licensed_features(custom_fields: true)
+    end
+
+    it 'updates the custom fields' do
+      existing_text_value = create(
+        :work_item_text_field_value, work_item: work_item, custom_field: text_field, value: 'old text'
+      )
+
+      post_graphql_mutation(mutation, current_user: reporter)
+
+      expect(response).to have_gitlab_http_status(:success)
+
+      expect(existing_text_value.reload.value).to eq('some text')
+      expect(WorkItems::SelectFieldValue.last).to have_attributes(
+        work_item_id: work_item.id,
+        custom_field_id: select_field.id,
+        custom_field_select_option_id: select_option_1.id
+      )
+    end
+  end
+
   def mutation_for(item)
     graphql_mutation(:workItemUpdate, input.merge('id' => item.to_global_id.to_s), fields)
   end

ee/spec/services/work_items/callbacks/custom_fields_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe WorkItems::Callbacks::CustomFields, feature_category: :team_planning do
+  include_context 'with group configured with custom fields'
+
+  let_it_be(:project) { create(:project, group: group) }
+  let_it_be(:work_item) { create(:work_item, work_item_type: issue_type, project: project) }
+
+  let_it_be(:current_user) { create(:user, reporter_of: group) }
+
+  let(:callback) { described_class.new(issuable: work_item, current_user: current_user, params: params) }
+
+  let(:params) do
+    [
+      { custom_field_id: text_field.id, text_value: 'some text' },
+      { custom_field_id: number_field.id, number_value: 100 },
+      { custom_field_id: select_field.id, selected_option_ids: [select_option_1.id] },
+      { custom_field_id: multi_select_field.id, selected_option_ids: [
+        multi_select_option_1.id, multi_select_option_2.id
+      ] }
+    ]
+  end
+
+  before do
+    stub_licensed_features(custom_fields: true)
+  end
+
+  describe '#after_save' do
+    subject(:after_save_callback) { callback.after_save }
+
+    it 'sets the custom field values for the work item' do
+      after_save_callback
+
+      expect(custom_field_values_for(text_field).first.value).to eq('some text')
+      expect(custom_field_values_for(number_field).first.value).to eq(100)
+      expect(custom_field_values_for(select_field).first.custom_field_select_option_id).to eq(select_option_1.id)
+      expect(custom_field_values_for(multi_select_field)).to contain_exactly(
+        have_attributes(custom_field_select_option_id: multi_select_option_1.id),
+        have_attributes(custom_field_select_option_id: multi_select_option_2.id)
+      )
+    end
+
+    context 'when there are existing values' do
+      before do
+        create(:work_item_text_field_value, work_item: work_item, custom_field: text_field, value: 'existing text')
+        create(:work_item_number_field_value, work_item: work_item, custom_field: number_field, value: 50)
+      end
+
+      let(:params) do
+        [
+          { custom_field_id: number_field.id, number_value: 100 }
+        ]
+      end
+
+      it 'updates the given fields and leaves others as-is' do
+        after_save_callback
+
+        expect(custom_field_values_for(text_field).first.value).to eq('existing text')
+        expect(custom_field_values_for(number_field).first.value).to eq(100)
+      end
+    end
+
+    context 'when custom_fields_feature is disabled' do
+      before do
+        stub_feature_flags(custom_fields_feature: false)
+      end
+
+      it 'does not set any custom field values' do
+        after_save_callback
+
+        expect(custom_field_values_for(text_field)).to be_empty
+        expect(custom_field_values_for(number_field)).to be_empty
+        expect(custom_field_values_for(select_field)).to be_empty
+        expect(custom_field_values_for(multi_select_field)).to be_empty
+      end
+    end
+
+    context 'when user does not have access' do
+      let(:current_user) { create(:user) }
+
+      it 'does not set any custom field values' do
+        after_save_callback
+
+        expect(custom_field_values_for(text_field)).to be_empty
+        expect(custom_field_values_for(number_field)).to be_empty
+        expect(custom_field_values_for(select_field)).to be_empty
+        expect(custom_field_values_for(multi_select_field)).to be_empty
+      end
+    end
+
+    context 'when custom field ID is invalid' do
+      let(:params) do
+        [
+          { custom_field_id: non_existing_record_id, text_value: 'some text' }
+        ]
+      end
+
+      it 'raises an error' do
+        expect { after_save_callback }.to raise_error(Issuable::Callbacks::Base::Error, /Invalid custom field ID/)
+      end
+    end
+
+    context 'when there are validation errors' do
+      let(:params) do
+        [
+          { custom_field_id: number_field.id, number_value: 'some text' }
+        ]
+      end
+
+      it 'raises an error' do
+        expect do
+          after_save_callback
+        end.to raise_error(Issuable::Callbacks::Base::Error, 'Validation failed: Value is not a number')
+      end
+    end
+
+    context 'with unsupported custom field type' do
+      let(:invalid_field) do
+        create(:custom_field, namespace: group, work_item_types: [issue_type]).tap do |f|
+          f.update_column(:field_type, -1)
+        end
+      end
+
+      let(:params) do
+        [
+          { custom_field_id: invalid_field.id, text_value: 'some text' }
+        ]
+      end
+
+      it 'raises an error' do
+        expect do
+          after_save_callback
+        end.to raise_error(Issuable::Callbacks::Base::Error, /\AUnsupported field type/)
+      end
+    end
+  end
+
+  def custom_field_values_for(custom_field)
+    value_class = if custom_field.field_type_text?
+                    WorkItems::TextFieldValue
+                  elsif custom_field.field_type_number?
+                    WorkItems::NumberFieldValue
+                  elsif custom_field.field_type_select?
+                    WorkItems::SelectFieldValue
+                  end
+
+    value_class.where(work_item: work_item, custom_field: custom_field)
+  end
+end