diff --git a/assets/javascripts/discourse/components/ai-composer-helper-menu.gjs b/assets/javascripts/discourse/components/ai-composer-helper-menu.gjs
index f5c4a6c3..c9f26f0a 100644
--- a/assets/javascripts/discourse/components/ai-composer-helper-menu.gjs
+++ b/assets/javascripts/discourse/components/ai-composer-helper-menu.gjs
@@ -1,6 +1,5 @@
 import Component from "@glimmer/component";
 import { tracked } from "@glimmer/tracking";
-import { fn } from "@ember/helper";
 import { action } from "@ember/object";
 import { service } from "@ember/service";
 import { modifier } from "ember-modifier";
@@ -17,6 +16,7 @@ import ModalDiffModal from "../components/modal/diff-modal";
 import ThumbnailSuggestion from "../components/modal/thumbnail-suggestions";
 
 export default class AiComposerHelperMenu extends Component {
+  @service modal;
   @service siteSettings;
   @service aiComposerHelper;
   @service currentUser;
@@ -29,7 +29,6 @@ export default class AiComposerHelperMenu extends Component {
   @tracked lastUsedOption = null;
   @tracked thumbnailSuggestions = null;
   @tracked showThumbnailModal = false;
-  @tracked showDiffModal = false;
   @tracked lastSelectionRange = null;
   MENU_STATES = this.aiComposerHelper.MENU_STATES;
   prompts = [];
@@ -100,7 +99,16 @@ export default class AiComposerHelperMenu extends Component {
       {
         icon: "exchange-alt",
         label: "discourse_ai.ai_helper.context_menu.view_changes",
-        action: () => (this.showDiffModal = true),
+        action: () =>
+          this.modal.show(ModalDiffModal, {
+            model: {
+              diff: this.diff,
+              oldValue: this.initialValue,
+              newValue: this.newSelectedText,
+              revert: this.undoAiAction,
+              confirm: () => this.updateMenuState(this.MENU_STATES.resets),
+            },
+          }),
         classes: "view-changes",
       },
       {
@@ -202,8 +210,12 @@ export default class AiComposerHelperMenu extends Component {
       if (option.name === "illustrate_post") {
         this._toggleLoadingState(false);
         this.closeMenu();
-        this.showThumbnailModal = true;
         this.thumbnailSuggestions = data.thumbnails;
+        this.modal.show(ThumbnailSuggestion, {
+          model: {
+            thumbnails: this.thumbnailSuggestions,
+          },
+        });
       } else {
         this._updateSuggestedByAi(data);
       }
@@ -326,26 +338,5 @@ export default class AiComposerHelperMenu extends Component {
         />
       {{/if}}
     </div>
-
-    {{#if this.showDiffModal}}
-      <ModalDiffModal
-        @confirm={{fn
-          (mut this.aiComposerHelper.menuState)
-          this.MENU_STATES.resets
-        }}
-        @diff={{this.diff}}
-        @oldValue={{this.initialValue}}
-        @newValue={{this.newSelectedText}}
-        @revert={{this.undoAiAction}}
-        @closeModal={{fn (mut this.showDiffModal) false}}
-      />
-    {{/if}}
-
-    {{#if this.showThumbnailModal}}
-      <ThumbnailSuggestion
-        @thumbnails={{this.thumbnailSuggestions}}
-        @closeModal={{fn (mut this.showThumbnailModal) false}}
-      />
-    {{/if}}
   </template>
 }
diff --git a/assets/javascripts/discourse/components/modal/diff-modal.gjs b/assets/javascripts/discourse/components/modal/diff-modal.gjs
index 8edf0cfa..66cba0f0 100644
--- a/assets/javascripts/discourse/components/modal/diff-modal.gjs
+++ b/assets/javascripts/discourse/components/modal/diff-modal.gjs
@@ -9,13 +9,13 @@ export default class ModalDiffModal extends Component {
   @action
   triggerConfirmChanges() {
     this.args.closeModal();
-    this.args.confirm();
+    this.args.model.confirm();
   }
 
   @action
   triggerRevertChanges() {
+    this.args.model.revert();
     this.args.closeModal();
-    this.args.revert();
   }
 
   <template>
@@ -25,15 +25,15 @@ export default class ModalDiffModal extends Component {
       @closeModal={{@closeModal}}
     >
       <:body>
-        {{#if @diff}}
-          {{htmlSafe @diff}}
+        {{#if @model.diff}}
+          {{htmlSafe @model.diff}}
         {{else}}
           <div class="composer-ai-helper-modal__old-value">
-            {{@oldValue}}
+            {{@model.oldValue}}
           </div>
 
           <div class="composer-ai-helper-modal__new-value">
-            {{@newValue}}
+            {{@model.newValue}}
           </div>
         {{/if}}
       </:body>
@@ -45,7 +45,7 @@ export default class ModalDiffModal extends Component {
           @label="discourse_ai.ai_helper.context_menu.confirm"
         />
         <DButton
-          class="btn-flat"
+          class="btn-flat revert"
           @action={{this.triggerRevertChanges}}
           @label="discourse_ai.ai_helper.context_menu.revert"
         />
diff --git a/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs b/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs
index b56a37d8..09ef82e2 100644
--- a/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs
+++ b/assets/javascripts/discourse/components/modal/thumbnail-suggestions.gjs
@@ -53,7 +53,7 @@ export default class ThumbnailSuggestions extends Component {
     >
       <:body>
         <div class="ai-thumbnail-suggestions">
-          {{#each @thumbnails as |thumbnail|}}
+          {{#each @model.thumbnails as |thumbnail|}}
             <ThumbnailSuggestionItem
               @thumbnail={{thumbnail}}
               @addSelection={{this.addSelection}}
diff --git a/spec/system/ai_helper/ai_composer_helper_spec.rb b/spec/system/ai_helper/ai_composer_helper_spec.rb
index fa5695a2..c4c4953e 100644
--- a/spec/system/ai_helper/ai_composer_helper_spec.rb
+++ b/spec/system/ai_helper/ai_composer_helper_spec.rb
@@ -242,6 +242,23 @@ RSpec.describe "AI Composer helper", type: :system, js: true do
           expect(ai_helper_context_menu).to have_no_context_menu
         end
       end
+
+      it "reverts the changes when revert button is pressed in the modal" do
+        trigger_context_menu(spanish_input)
+        ai_helper_context_menu.click_ai_button
+
+        DiscourseAi::Completions::Llm.with_prepared_responses([input]) do
+          ai_helper_context_menu.select_helper_model(mode)
+
+          wait_for { composer.composer_input.value == input }
+
+          ai_helper_context_menu.click_view_changes_button
+          expect(diff_modal).to be_visible
+          diff_modal.revert_changes
+          expect(ai_helper_context_menu).to have_no_context_menu
+          expect(composer.composer_input.value).to eq(spanish_input)
+        end
+      end
     end
 
     context "when using the proofreading mode" do
diff --git a/spec/system/page_objects/modals/diff_modal.rb b/spec/system/page_objects/modals/diff_modal.rb
index 4caa03a9..9082ddaa 100644
--- a/spec/system/page_objects/modals/diff_modal.rb
+++ b/spec/system/page_objects/modals/diff_modal.rb
@@ -11,6 +11,10 @@ module PageObjects
         find(".d-modal__footer button.confirm", wait: 5).click
       end
 
+      def revert_changes
+        find(".d-modal__footer button.revert", wait: 5).click
+      end
+
       def old_value
         find(".composer-ai-helper-modal__old-value").text
       end