Make model execution features support immediate enable/disable

Add feature state change observer for notifying consumers, and deprecate the enabling on browser restart notifications.

Removes the restart toast from settings AI page, since its no longer needed.

Also adds enterprise policy pref observer for model execution controller to observer immediate changes.

Demo of settings change immediately affecting feature state:
http://go/scrcast/NDYzMTcwMDEzMzkwNDM4NHwyMzAyNDZlYS1lNA

Bug: b/320851152,  b/321073908, b/323569816
Change-Id: I7c4c4fe3bb68527df8e22247dbfc3aceb5b020e4
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5247445
Commit-Queue: Raj T <rajendrant@chromium.org>
Reviewed-by: Tibor Goldschwendt <tiborg@chromium.org>
Reviewed-by: Demetrios Papadopoulos <dpapad@chromium.org>
Reviewed-by: Carlos Knippschild <carlosk@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1255856}

chrome/browser/compose/chrome_compose_client.cc

@@ -568,6 +568,9 @@ void ChromeComposeClient::PrimaryPageChanged(content::Page& page) {
 
 void ChromeComposeClient::OnWebContentsFocused(
     content::RenderWidgetHost* render_widget_host) {
+  if (!compose_enabling_->IsEnabledForProfile(profile_)) {
+    return;
+  }
   ComposeSession* active_session = GetSessionForActiveComposeField();
   if (open_settings_requested_) {
     open_settings_requested_ = false;

chrome/browser/compose/compose_enabling_browsertest.cc

@@ -80,12 +80,9 @@ IN_PROC_BROWSER_TEST_F(ComposeEnablingBrowserTest,
               MODEL_EXECUTION_FEATURE_COMPOSE),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
 
-  // Checks that Compose is still disabled.
-  // TODO(b/321073908): Update this test once enabling takes immediate effect.
-  EXPECT_EQ(base::unexpected(
-                compose::ComposeShowStatus::kUserNotAllowedByOptimizationGuide),
-            GetComposeEnabling().IsEnabled());
-  EXPECT_FALSE(GetOptimizationGuide()->ShouldFeatureBeCurrentlyEnabledForUser(
+  // Checks that Compose is immediately enabled.
+  EXPECT_EQ(base::ok(), GetComposeEnabling().IsEnabled());
+  EXPECT_TRUE(GetOptimizationGuide()->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
           MODEL_EXECUTION_FEATURE_COMPOSE));
 }

chrome/browser/optimization_guide/model_execution/model_execution_browsertest.cc

@@ -627,15 +627,15 @@ class ModelExecutionEnterprisePolicyBrowserTest
 
 IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
                        DisableThenEnable) {
+  EnableSignin();
+
   auto* prefs = browser()->profile()->GetPrefs();
   prefs->SetInteger(
       prefs::GetSettingEnabledPrefName(
           optimization_guide::proto::ModelExecutionFeature::
               MODEL_EXECUTION_FEATURE_TAB_ORGANIZATION),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
-  GetOptGuideKeyedService()->SimulateBrowserRestartForControllerTesting();
-
-  EnableSignin();
+  base::RunLoop().RunUntilIdle();
 
   // Default policy value allows the feature.
   EXPECT_TRUE(IsSettingVisible(
@@ -679,6 +679,11 @@ IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
           model_execution::prefs::ModelExecutionEnterprisePolicyValue::kAllow)),
       nullptr);
   policy_provider_.UpdateChromePolicy(policies);
+  prefs->SetInteger(
+      prefs::GetSettingEnabledPrefName(
+          optimization_guide::proto::ModelExecutionFeature::
+              MODEL_EXECUTION_FEATURE_TAB_ORGANIZATION),
+      static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
   base::RunLoop().RunUntilIdle();
   EXPECT_TRUE(IsSettingVisible(
       proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TAB_ORGANIZATION));
@@ -694,15 +699,15 @@ IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
 
 IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
                        DisableThenEnableCompose) {
+  EnableSignin();
+
   auto* prefs = browser()->profile()->GetPrefs();
   prefs->SetInteger(
       prefs::GetSettingEnabledPrefName(
           optimization_guide::proto::ModelExecutionFeature::
               MODEL_EXECUTION_FEATURE_COMPOSE),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
-  GetOptGuideKeyedService()->SimulateBrowserRestartForControllerTesting();
-
-  EnableSignin();
+  base::RunLoop().RunUntilIdle();
 
   // Default policy value allows the feature.
   EXPECT_TRUE(IsSettingVisible(
@@ -733,6 +738,11 @@ IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
           model_execution::prefs::ModelExecutionEnterprisePolicyValue::kAllow)),
       nullptr);
   policy_provider_.UpdateChromePolicy(policies);
+  prefs->SetInteger(
+      prefs::GetSettingEnabledPrefName(
+          optimization_guide::proto::ModelExecutionFeature::
+              MODEL_EXECUTION_FEATURE_COMPOSE),
+      static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
   base::RunLoop().RunUntilIdle();
   EXPECT_TRUE(IsSettingVisible(
       proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_COMPOSE));
@@ -742,15 +752,15 @@ IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
 
 IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
                        DisableThenEnableWallpaperSearch) {
+  EnableSignin();
+
   auto* prefs = browser()->profile()->GetPrefs();
   prefs->SetInteger(
       prefs::GetSettingEnabledPrefName(
           optimization_guide::proto::ModelExecutionFeature::
               MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
-  GetOptGuideKeyedService()->SimulateBrowserRestartForControllerTesting();
-
-  EnableSignin();
+  base::RunLoop().RunUntilIdle();
 
   // Default policy value allows the feature.
   EXPECT_TRUE(IsSettingVisible(
@@ -782,6 +792,11 @@ IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
           model_execution::prefs::ModelExecutionEnterprisePolicyValue::kAllow)),
       nullptr);
   policy_provider_.UpdateChromePolicy(policies);
+  prefs->SetInteger(
+      prefs::GetSettingEnabledPrefName(
+          optimization_guide::proto::ModelExecutionFeature::
+              MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH),
+      static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
   base::RunLoop().RunUntilIdle();
   EXPECT_TRUE(IsSettingVisible(
       proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
@@ -791,15 +806,15 @@ IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
 
 IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
                        EnableWithoutLogging) {
+  EnableSignin();
+
   auto* prefs = browser()->profile()->GetPrefs();
   prefs->SetInteger(
       prefs::GetSettingEnabledPrefName(
           optimization_guide::proto::ModelExecutionFeature::
               MODEL_EXECUTION_FEATURE_TAB_ORGANIZATION),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
-  GetOptGuideKeyedService()->SimulateBrowserRestartForControllerTesting();
-
-  EnableSignin();
+  base::RunLoop().RunUntilIdle();
 
   // EnableWithoutLogging via the enterprise policy.
   policy::PolicyMap policies;
@@ -811,6 +826,11 @@ IN_PROC_BROWSER_TEST_F(ModelExecutionEnterprisePolicyBrowserTest,
                        kAllowWithoutLogging)),
                nullptr);
   policy_provider_.UpdateChromePolicy(policies);
+  prefs->SetInteger(
+      prefs::GetSettingEnabledPrefName(
+          optimization_guide::proto::ModelExecutionFeature::
+              MODEL_EXECUTION_FEATURE_TAB_ORGANIZATION),
+      static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
   base::RunLoop().RunUntilIdle();
   EXPECT_TRUE(IsSettingVisible(
       proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TAB_ORGANIZATION));

chrome/browser/optimization_guide/optimization_guide_keyed_service.cc

@@ -250,16 +250,6 @@ OptimizationGuideKeyedService::MaybeCreatePushNotificationManager(
   return nullptr;
 }
 
-void OptimizationGuideKeyedService::
-    SimulateBrowserRestartForControllerTesting() {
-  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
-  if (!model_execution_features_controller_) {
-    return;
-  }
-  model_execution_features_controller_
-      ->SimulateBrowserRestartForTesting();  // IN-TEST
-}
-
 // static
 // We're using a weakptr here for testing purposes. We need to allow
 // OnDeviceModelComponentStateManager to be destroyed along with a test harness.

chrome/browser/optimization_guide/optimization_guide_keyed_service.h

@@ -176,10 +176,6 @@ class OptimizationGuideKeyedService
     return optimization_guide_logger_.get();
   }
 
-  // Simulates browser restart. Useful for testing controller functionality
-  // where some of the settings change take effect on browser restart.
-  void SimulateBrowserRestartForControllerTesting();
-
  private:
   friend class BrowserView;
   friend class ChromeBrowserMainExtraPartsOptimizationGuide;

chrome/browser/optimization_guide/optimization_guide_keyed_service_browsertest.cc

@@ -793,11 +793,14 @@ class TestSettingsEnabledObserver
   explicit TestSettingsEnabledObserver(
       optimization_guide::proto::ModelExecutionFeature feature)
       : SettingsEnabledObserver(feature) {}
-  void PrepareToEnableOnRestart() override {
-    ++count_received_prepare_to_enable_on_restart_notifications_;
+  void PrepareToEnableOnRestart() override {}
+  void OnChangeInFeatureCurrentlyEnabledState(bool is_now_enabled) override {
+    count_feature_enabled_state_changes_++;
+    is_currently_enabled_ = is_now_enabled;
   }
 
-  int count_received_prepare_to_enable_on_restart_notifications_ = 0;
+  int count_feature_enabled_state_changes_ = 0;
+  bool is_currently_enabled_ = false;
 };
 
 IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
@@ -874,15 +877,6 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
               MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
 
-  EXPECT_FALSE(
-      IsSettingVisible(optimization_guide::proto::ModelExecutionFeature::
-                           MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-
-  OptimizationGuideKeyedService* ogks =
-      OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile());
-
-  ogks->SimulateBrowserRestartForControllerTesting();
-
   // Restarting the browser should cause wallpaper setting to be visible since
   // the feature is enabled.
   EXPECT_TRUE(
@@ -904,12 +898,6 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
       static_cast<int>(
           optimization_guide::prefs::FeatureOptInState::kDisabled));
 
-  EXPECT_TRUE(
-      IsSettingVisible(optimization_guide::proto::ModelExecutionFeature::
-                           MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-
-  ogks->SimulateBrowserRestartForControllerTesting();
-
   // Restarting the browser should cause wallpaper setting to be not visible
   // since the feature is no-longer enabled.
   EXPECT_FALSE(
@@ -960,16 +948,9 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
           optimization_guide::proto::ModelExecutionFeature::
               MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
-  EXPECT_EQ(1, wallpaper_search_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-  EXPECT_EQ(0, compose_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-
-  EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
-      optimization_guide::proto::ModelExecutionFeature::
-          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-
-  ogks->SimulateBrowserRestartForControllerTesting();
+  EXPECT_EQ(1, wallpaper_search_observer.count_feature_enabled_state_changes_);
+  EXPECT_TRUE(wallpaper_search_observer.is_currently_enabled_);
+  EXPECT_EQ(0, compose_observer.count_feature_enabled_state_changes_);
 
   EXPECT_TRUE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
@@ -998,6 +979,9 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
   EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
           MODEL_EXECUTION_FEATURE_COMPOSE));
+
+  EXPECT_EQ(2, wallpaper_search_observer.count_feature_enabled_state_changes_);
+  EXPECT_FALSE(wallpaper_search_observer.is_currently_enabled_);
 #endif
 }
 
@@ -1038,16 +1022,9 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
           optimization_guide::proto::ModelExecutionFeature::
               MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
-  EXPECT_EQ(1, wallpaper_search_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-  EXPECT_EQ(0, compose_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-
-  EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
-      optimization_guide::proto::ModelExecutionFeature::
-          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-
-  ogks->SimulateBrowserRestartForControllerTesting();
+  EXPECT_EQ(1, wallpaper_search_observer.count_feature_enabled_state_changes_);
+  EXPECT_TRUE(wallpaper_search_observer.is_currently_enabled_);
+  EXPECT_EQ(0, compose_observer.count_feature_enabled_state_changes_);
 
   EXPECT_TRUE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
@@ -1067,16 +1044,9 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
               MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH),
       static_cast<int>(
           optimization_guide::prefs::FeatureOptInState::kDisabled));
-  EXPECT_EQ(1, wallpaper_search_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-  EXPECT_EQ(0, compose_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-
-  EXPECT_TRUE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
-      optimization_guide::proto::ModelExecutionFeature::
-          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-
-  ogks->SimulateBrowserRestartForControllerTesting();
+  EXPECT_EQ(2, wallpaper_search_observer.count_feature_enabled_state_changes_);
+  EXPECT_FALSE(wallpaper_search_observer.is_currently_enabled_);
+  EXPECT_EQ(0, compose_observer.count_feature_enabled_state_changes_);
 
   EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
@@ -1134,19 +1104,11 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
   // Visibility of tab organizer feature is enabled via finch. Only tab
   // organizer feature should be enabled.
-  EXPECT_EQ(0, wallpaper_search_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-  EXPECT_EQ(1, compose_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-  EXPECT_EQ(
-      1,
-      tab_observer.count_received_prepare_to_enable_on_restart_notifications_);
-
-  EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
-      optimization_guide::proto::ModelExecutionFeature::
-          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-
-  ogks->SimulateBrowserRestartForControllerTesting();
+  EXPECT_EQ(0, wallpaper_search_observer.count_feature_enabled_state_changes_);
+  EXPECT_EQ(1, compose_observer.count_feature_enabled_state_changes_);
+  EXPECT_TRUE(compose_observer.is_currently_enabled_);
+  EXPECT_EQ(1, tab_observer.count_feature_enabled_state_changes_);
+  EXPECT_TRUE(tab_observer.is_currently_enabled_);
 
   EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
@@ -1166,19 +1128,13 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
       optimization_guide::prefs::kModelExecutionMainToggleSettingState,
       static_cast<int>(
           optimization_guide::prefs::FeatureOptInState::kDisabled));
-  EXPECT_EQ(0, wallpaper_search_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-  EXPECT_EQ(1, compose_observer
-                   .count_received_prepare_to_enable_on_restart_notifications_);
-  EXPECT_EQ(
-      1,
-      tab_observer.count_received_prepare_to_enable_on_restart_notifications_);
+  base::RunLoop().RunUntilIdle();
 
-  EXPECT_TRUE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
-      optimization_guide::proto::ModelExecutionFeature::
-          MODEL_EXECUTION_FEATURE_TAB_ORGANIZATION));
-
-  ogks->SimulateBrowserRestartForControllerTesting();
+  EXPECT_EQ(1, wallpaper_search_observer.count_feature_enabled_state_changes_);
+  EXPECT_EQ(2, compose_observer.count_feature_enabled_state_changes_);
+  EXPECT_FALSE(compose_observer.is_currently_enabled_);
+  EXPECT_EQ(2, tab_observer.count_feature_enabled_state_changes_);
+  EXPECT_FALSE(tab_observer.is_currently_enabled_);
 
   EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
@@ -1197,6 +1153,8 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
 // profiles.
 IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
                        SettingsVisibilityIncognito) {
+  EnableSignIn();
+
   // Set up incognito browser and incognito OptimizationGuideKeyedService
   // consumer.
   Browser* otr_browser = CreateIncognitoBrowser(browser()->profile());
@@ -1302,16 +1260,6 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceBrowserTest,
               MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
 
-  EXPECT_FALSE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
-      optimization_guide::proto::ModelExecutionFeature::
-          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-  EXPECT_FALSE(guest_ogks->ShouldFeatureBeCurrentlyEnabledForUser(
-      optimization_guide::proto::ModelExecutionFeature::
-          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
-
-  ogks->SimulateBrowserRestartForControllerTesting();
-  guest_ogks->SimulateBrowserRestartForControllerTesting();
-
   EXPECT_TRUE(ogks->ShouldFeatureBeCurrentlyEnabledForUser(
       optimization_guide::proto::ModelExecutionFeature::
           MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
@@ -1561,7 +1509,7 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceEnterpriseBrowserTest,
   prefs->SetInteger(
       optimization_guide::prefs::GetSettingEnabledPrefName(compose_feature),
       static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
-  ogks->SimulateBrowserRestartForControllerTesting();
+  base::RunLoop().RunUntilIdle();
 
   policy::PolicyMap policies;
 
@@ -1616,6 +1564,9 @@ IN_PROC_BROWSER_TEST_F(OptimizationGuideKeyedServiceEnterpriseBrowserTest,
                        ModelExecutionEnterprisePolicyValue::kAllow)),
                nullptr);
   policy_provider_.UpdateChromePolicy(policies);
+  prefs->SetInteger(
+      optimization_guide::prefs::GetSettingEnabledPrefName(compose_feature),
+      static_cast<int>(optimization_guide::prefs::FeatureOptInState::kEnabled));
   base::RunLoop().RunUntilIdle();
 
   EXPECT_TRUE(ogks->ShouldFeatureBeCurrentlyAllowedForLogging(compose_feature));

chrome/browser/resources/settings/ai_page/ai_page.html

@@ -13,8 +13,7 @@
     numeric-unchecked-values="[[numericUncheckedValues_]]"
     numeric-checked-value="[[featureOptInStateEnum_.ENABLED]]"
     label="$i18n{aiPageMainLabel}"
-    sub-label="$i18n{aiPageMainSublabel}"
-    on-settings-boolean-control-change="onToggleChange_">
+    sub-label="$i18n{aiPageMainSublabel}">
 </settings-toggle-button>
 
 <iron-collapse opened="[[isExpanded_(
@@ -26,8 +25,7 @@
         numeric-unchecked-values="[[numericUncheckedValues_]]"
         numeric-checked-value="[[featureOptInStateEnum_.ENABLED]]"
         label="$i18n{aiComposeLabel}"
-        sub-label="$i18n{aiComposeSublabel}"
-        on-settings-boolean-control-change="onToggleChange_">
+        sub-label="$i18n{aiComposeSublabel}">
     </settings-toggle-button>
     <settings-toggle-button
         class$="[[getTabOrganizationHrCssClass_(showComposeControl_)]]"
@@ -36,8 +34,7 @@
         numeric-unchecked-values="[[numericUncheckedValues_]]"
         numeric-checked-value="[[featureOptInStateEnum_.ENABLED]]"
         label="$i18n{experimentalAdvancedFeature2Label}"
-        sub-label="$i18n{experimentalAdvancedFeature2Sublabel}"
-        on-settings-boolean-control-change="onToggleChange_">
+        sub-label="$i18n{experimentalAdvancedFeature2Sublabel}">
     </settings-toggle-button>
     <settings-toggle-button
         class$="[[getWallpaperSearchHrCssClass_(
@@ -47,20 +44,8 @@
         numeric-unchecked-values="[[numericUncheckedValues_]]"
         numeric-checked-value="[[featureOptInStateEnum_.ENABLED]]"
         label="$i18n{experimentalAdvancedFeature3Label}"
-        sub-label="$i18n{experimentalAdvancedFeature3Sublabel}"
-        on-settings-boolean-control-change="onToggleChange_">
+        sub-label="$i18n{experimentalAdvancedFeature3Sublabel}">
     </settings-toggle-button>
   </div>
 </iron-collapse>
 
-<cr-toast id="toast">
-  <div>$i18n{restartToApplyChanges}</div>
-  <cr-button on-click="onRestartClick_">$i18n{restart}</cr-button>
-</cr-toast>
-
-<if expr="not chromeos_ash">
-<template is="dom-if" if="[[shouldShowRelaunchDialog]]" restamp>
-  <relaunch-confirmation-dialog restart-type="[[restartTypeEnum.RESTART]]"
-      on-close="onRelaunchDialogClose"></relaunch-confirmation-dialog>
-</template>
-</if>

chrome/browser/resources/settings/ai_page/ai_page.ts

@@ -3,20 +3,12 @@
 // found in the LICENSE file.
 
 import '../controls/settings_toggle_button.js';
-import 'chrome://resources/cr_elements/cr_button/cr_button.js';
-import 'chrome://resources/cr_elements/cr_toast/cr_toast.js';
 import 'chrome://resources/polymer/v3_0/iron-collapse/iron-collapse.js';
-// <if expr="not chromeos_ash">
-import '../relaunch_confirmation_dialog.js';
-
-// </if>
 
 import {PrefsMixin} from 'chrome://resources/cr_components/settings_prefs/prefs_mixin.js';
-import type {CrToastElement} from 'chrome://resources/cr_elements/cr_toast/cr_toast.js';
 import {PolymerElement} from 'chrome://resources/polymer/v3_0/polymer/polymer_bundled.min.js';
 
 import {loadTimeData} from '../i18n_setup.js';
-import {RelaunchMixin, RestartType} from '../relaunch_mixin.js';
 
 import {getTemplate} from './ai_page.html.js';
 
@@ -37,13 +29,7 @@ export enum SettingsAiPageFeaturePrefName {
   WALLPAPER_SEARCH = 'optimization_guide.wallpaper_search_setting_state',
 }
 
-export interface SettingsAiPageElement {
-  $: {
-    toast: CrToastElement,
-  };
-}
-
-const SettingsAiPageElementBase = RelaunchMixin(PrefsMixin(PolymerElement));
+const SettingsAiPageElementBase = PrefsMixin(PolymerElement);
 
 export class SettingsAiPageElement extends SettingsAiPageElementBase {
   static get is() {
@@ -94,15 +80,6 @@ export class SettingsAiPageElement extends SettingsAiPageElementBase {
   private showWallpaperSearchControl_: boolean;
   private numericUncheckedValues_: FeatureOptInState[];
 
-  private onToggleChange_() {
-    this.$.toast.show();
-  }
-
-  private onRestartClick_(e: Event) {
-    e.stopPropagation();
-    this.performRestart(RestartType.RESTART);
-  }
-
   private isExpanded_(): boolean {
     return this.getPref(SettingsAiPageFeaturePrefName.MAIN).value ===
         FeatureOptInState.ENABLED;

chrome/browser/search/background/wallpaper_search/wallpaper_search_service_browsertest.cc

@@ -7,6 +7,7 @@
 #include "base/feature_list.h"
 #include "base/test/scoped_feature_list.h"
 #include "chrome/browser/optimization_guide/browser_test_util.h"
+#include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/search/background/wallpaper_search/wallpaper_search_service_factory.h"
 #include "chrome/browser/signin/identity_manager_factory.h"
@@ -51,7 +52,7 @@ class WallpaperSearchServiceBrowserTest : public InProcessBrowserTest {
 #if !BUILDFLAG(IS_CHROMEOS_ASH)
 // PRE_ simulates a browser restart.
 IN_PROC_BROWSER_TEST_F(WallpaperSearchServiceBrowserTest,
-                       PRE_EnablingWallpaperSearchEnablesGM3) {
+                       PRE_EnablingWallpaperSearchEnables) {
   optimization_guide::EnableSigninAndModelExecutionCapability(
       browser()->profile());
 
@@ -65,10 +66,13 @@ IN_PROC_BROWSER_TEST_F(WallpaperSearchServiceBrowserTest,
 }
 
 IN_PROC_BROWSER_TEST_F(WallpaperSearchServiceBrowserTest,
-                       EnablingWallpaperSearchEnablesGM3) {
-  EXPECT_TRUE(base::FeatureList::IsEnabled(features::kChromeRefresh2023));
-  EXPECT_TRUE(base::FeatureList::IsEnabled(features::kChromeWebuiRefresh2023));
-  EXPECT_TRUE(features::IsChromeWebuiRefresh2023());
+                       EnablingWallpaperSearchEnables) {
+  // Wallpaper search feature should be enabled.
+  auto* keyed_service =
+      OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile());
+  EXPECT_TRUE(keyed_service->ShouldFeatureBeCurrentlyEnabledForUser(
+      optimization_guide::proto::ModelExecutionFeature::
+          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
 }
 #endif  // !BUILDFLAG(IS_CHROMEOS_ASH)
 
@@ -88,7 +92,7 @@ INSTANTIATE_TEST_SUITE_P(All,
                                             ::testing::Bool()));
 
 IN_PROC_BROWSER_TEST_P(WallpaperSearchServiceBrowserChromeAshTest,
-                       PRE_EnablingWallpaperSearchEnablesGM3) {
+                       PRE_EnablingWallpaperSearchEnables) {
   signin::MakePrimaryAccountAvailable(
       IdentityManagerFactory::GetForProfile(browser()->profile()),
       "test@example.com", signin::ConsentLevel::kSync);
@@ -110,18 +114,12 @@ IN_PROC_BROWSER_TEST_P(WallpaperSearchServiceBrowserChromeAshTest,
 }
 
 IN_PROC_BROWSER_TEST_P(WallpaperSearchServiceBrowserChromeAshTest,
-                       EnablingWallpaperSearchEnablesGM3) {
-#if !BUILDFLAG(IS_CHROMEOS_DEVICE)
-  // This test edits flags in chrome://flags, which doesn't work for
-  // Chrome-OS-on-Linux. It might cause parts of this test to flake.
-  // https://chromium.googlesource.com/chromium/src/+/HEAD/docs/chromeos_build_instructions.md#flags
-  if (!IsRunningOnChromeOS()) {
-    EXPECT_TRUE(base::FeatureList::IsEnabled(features::kChromeRefresh2023));
-  }
-#else
-  EXPECT_TRUE(base::FeatureList::IsEnabled(features::kChromeRefresh2023));
-  EXPECT_TRUE(base::FeatureList::IsEnabled(features::kChromeWebuiRefresh2023));
-  EXPECT_TRUE(features::IsChromeWebuiRefresh2023());
-#endif  // !BUILDFLAG(IS_CHROMEOS_DEVICE)
+                       EnablingWallpaperSearchEnables) {
+  // Wallpaper search feature should be enabled.
+  auto* keyed_service =
+      OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile());
+  EXPECT_TRUE(keyed_service->ShouldFeatureBeCurrentlyEnabledForUser(
+      optimization_guide::proto::ModelExecutionFeature::
+          MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH));
 }
 #endif  // BUILDFLAG(IS_CHROMEOS_ASH)

chrome/test/data/webui/settings/ai_page_test.ts

@@ -5,16 +5,14 @@
 'chrome://settings/settings.js';
 
 import type {SettingsToggleButtonElement, SettingsAiPageElement, SettingsPrefsElement} from 'chrome://settings/settings.js';
-import {LifetimeBrowserProxyImpl, SettingsAiPageFeaturePrefName as PrefName, CrSettingsPrefs, loadTimeData, FeatureOptInState} from 'chrome://settings/settings.js';
+import {SettingsAiPageFeaturePrefName as PrefName, CrSettingsPrefs, loadTimeData, FeatureOptInState} from 'chrome://settings/settings.js';
 
 import {assertEquals, assertTrue, assertFalse} from 'chrome://webui-test/chai_assert.js';
 import {isVisible} from 'chrome://webui-test/test_util.js';
-import {TestLifetimeBrowserProxy} from './test_lifetime_browser_proxy.js';
 
 suite('ExperimentalAdvancedPage', function() {
   let page: SettingsAiPageElement;
   let settingsPrefs: SettingsPrefsElement;
-  let lifetimeBrowserProxy: TestLifetimeBrowserProxy;
 
   suiteSetup(function() {
     settingsPrefs = document.createElement('settings-prefs');
@@ -22,9 +20,6 @@ suite('ExperimentalAdvancedPage', function() {
   });
 
   function createPage() {
-    lifetimeBrowserProxy = new TestLifetimeBrowserProxy();
-    LifetimeBrowserProxyImpl.setInstance(lifetimeBrowserProxy);
-
     document.body.innerHTML = window.trustedTypes!.emptyHTML;
     page = document.createElement('settings-ai-page');
     page.prefs = settingsPrefs.prefs;
@@ -34,7 +29,6 @@ suite('ExperimentalAdvancedPage', function() {
   // Test that interacting with the main toggle
   //  - updates the corresponding pref
   //  - updates the iron-collapse opened status
-  //  - shows the restart toast
   test('MainToggle', () => {
     createPage();
     page.setPrefValue(PrefName.MAIN, FeatureOptInState.NOT_INITIALIZED);
@@ -47,21 +41,18 @@ suite('ExperimentalAdvancedPage', function() {
     // Check NOT_INITIALIZED case.
     assertFalse(mainToggle.checked);
     assertFalse(collapse.opened);
-    assertFalse(page.$.toast.open);
 
     // Check ENABLED case.
     mainToggle.click();
     assertEquals(FeatureOptInState.ENABLED, page.getPref(PrefName.MAIN).value);
     assertTrue(mainToggle.checked);
     assertTrue(collapse.opened);
-    assertTrue(page.$.toast.open);
 
     // Check DISABLED case.
     mainToggle.click();
     assertEquals(FeatureOptInState.DISABLED, page.getPref(PrefName.MAIN).value);
     assertFalse(mainToggle.checked);
     assertFalse(collapse.opened);
-    assertTrue(page.$.toast.open);
   });
 
   test('FeatureTogglesVisibility', () => {
@@ -111,7 +102,6 @@ suite('ExperimentalAdvancedPage', function() {
         page.shadowRoot!.querySelectorAll<SettingsToggleButtonElement>(
             'iron-collapse settings-toggle-button');
     assertEquals(3, toggles.length);
-    assertFalse(page.$.toast.open);
 
     for (const toggle of toggles) {
       assertTrue(!!toggle.pref);
@@ -132,10 +122,6 @@ suite('ExperimentalAdvancedPage', function() {
     assertPrefs(
         FeatureOptInState.ENABLED, FeatureOptInState.NOT_INITIALIZED,
         FeatureOptInState.NOT_INITIALIZED);
-    assertTrue(page.$.toast.open);
-
-    // Manually hide toast to test the next toggle, normally stays open.
-    page.$.toast.hide();
 
     toggles[1]!.click();
     assertPrefs(
@@ -146,51 +132,22 @@ suite('ExperimentalAdvancedPage', function() {
     assertPrefs(
         FeatureOptInState.ENABLED, FeatureOptInState.ENABLED,
         FeatureOptInState.ENABLED);
-    assertTrue(page.$.toast.open);
-    page.$.toast.hide();
 
     // Check turning off toggles one by one.
     toggles[0]!.click();
     assertPrefs(
         FeatureOptInState.DISABLED, FeatureOptInState.ENABLED,
         FeatureOptInState.ENABLED);
-    assertTrue(page.$.toast.open);
-    page.$.toast.hide();
 
     toggles[1]!.click();
     assertPrefs(
         FeatureOptInState.DISABLED, FeatureOptInState.DISABLED,
         FeatureOptInState.ENABLED);
-    assertTrue(page.$.toast.open);
-    page.$.toast.hide();
 
     toggles[2]!.click();
     assertPrefs(
         FeatureOptInState.DISABLED, FeatureOptInState.DISABLED,
         FeatureOptInState.DISABLED);
-    assertTrue(page.$.toast.open);
-  });
-
-  test('RelaunchToast', () => {
-    loadTimeData.overrideValues({
-      showComposeControl: true,
-      showTabOrganizationControl: true,
-      showWallpaperSearchControl: true,
-    });
-    createPage();
-    const toggle = page.shadowRoot!.querySelector<SettingsToggleButtonElement>(
-        'iron-collapse settings-toggle-button');
-    assertTrue(!!toggle);
-    assertFalse(toggle.checked);
-    assertFalse(page.$.toast.open);
-
-    toggle.click();
-    assertTrue(page.$.toast.open);
-    const restartButton = page.$.toast.querySelector('cr-button');
-    assertTrue(!!restartButton);
-
-    restartButton.click();
-    return lifetimeBrowserProxy.whenCalled('restart');
   });
 
   test('FeatureTogglesSeparators', () => {

components/optimization_guide/core/model_execution/model_execution_features_controller.cc

@@ -166,36 +166,12 @@ bool ModelExecutionFeaturesController::ShouldFeatureBeCurrentlyEnabledForUser(
   DCHECK_CALLED_ON_VALID_THREAD(thread_checker_);
 
   ScopedFeatureCurrentlyEnabledHistogramRecorder metrics_recorder;
-
-  switch (GetCurrentUserValidityResult(feature)) {
-    case ModelExecutionFeaturesController::UserValidityResult::
-        kInvalidUnsignedUser:
-      metrics_recorder.SetResult(
-          feature, FeatureCurrentlyEnabledResult::kNotEnabledUnsignedUser);
-      return false;
-    case ModelExecutionFeaturesController::UserValidityResult::
-        kInvalidEnterprisePolicy:
-      metrics_recorder.SetResult(
-          feature, FeatureCurrentlyEnabledResult::kNotEnabledEnterprisePolicy);
-      return false;
-    case ModelExecutionFeaturesController::UserValidityResult::
-        kInvalidModelExecutionCapability:
-      metrics_recorder.SetResult(
-          feature,
-          FeatureCurrentlyEnabledResult::kNotEnabledModelExecutionCapability);
-      return false;
-    case ModelExecutionFeaturesController::UserValidityResult::kValid:
-      break;
-  }
-
-  bool result = features_enabled_at_startup_.find(static_cast<int>(feature)) !=
-                features_enabled_at_startup_.end();
-
+  bool is_enabled = GetPrefState(feature) == prefs::FeatureOptInState::kEnabled;
   metrics_recorder.SetResult(
-      feature, result ? FeatureCurrentlyEnabledResult::kEnabledAtStartup
-                      : FeatureCurrentlyEnabledResult::kNotEnabledAtStartup);
-
-  return result;
+      feature, is_enabled
+                   ? FeatureCurrentlyEnabledResult::kEnabledAtStartup
+                   : FeatureCurrentlyEnabledResult::kNotEnabledAtStartup);
+  return is_enabled;
 }
 
 bool ModelExecutionFeaturesController::
@@ -340,47 +316,52 @@ void ModelExecutionFeaturesController::OnFeatureSettingPrefChanged(
   DCHECK_CALLED_ON_VALID_THREAD(thread_checker_);
 
   auto pref_value = GetPrefState(feature);
+  bool is_enabled = ShouldFeatureBeCurrentlyEnabledForUser(feature);
+
+  // When the feature is enabled, check the user is valid to enable the
+  // feature.
+  CHECK(!is_enabled ||
+            GetCurrentUserValidityResult(feature) ==
+                ModelExecutionFeaturesController::UserValidityResult::kValid,
+        base::NotFatalUntil::M125);
+
   if (pref_value != prefs::FeatureOptInState::kNotInitialized) {
     base::UmaHistogramBoolean(
         base::StrCat(
             {"OptimizationGuide.ModelExecution.FeatureEnabledAtSettingsChange.",
              GetStringNameForModelExecutionFeature(feature)}),
-        pref_value == prefs::FeatureOptInState::kEnabled);
+        is_enabled);
   }
-
-  if (GetCurrentUserValidityResult(feature) !=
-      ModelExecutionFeaturesController::UserValidityResult::kValid) {
-    return;
-  }
-
   for (SettingsEnabledObserver& obs : observers_) {
     if (obs.feature() != feature) {
       continue;
     }
-    if (pref_value == prefs::FeatureOptInState::kEnabled) {
-      obs.PrepareToEnableOnRestart();
-    }
+    obs.OnChangeInFeatureCurrentlyEnabledState(is_enabled);
   }
 }
 
+void ModelExecutionFeaturesController::OnFeatureEnterprisePolicyPrefChanged(
+    proto::ModelExecutionFeature feature) {
+  DCHECK_CALLED_ON_VALID_THREAD(thread_checker_);
+  // When enterprise policy changes from allowed to disallowed, the feature
+  // settings prefs need to be cleared. This in turn triggers
+  // `OnFeatureSettingPrefChanged` to notify the observers of settings change,
+  // when the pref is reset.
+  ResetInvalidFeaturePrefs();
+}
+
 void ModelExecutionFeaturesController::InitializeFeatureSettings() {
-  features_enabled_at_startup_.clear();
   for (int i = 0; i < proto::ModelExecutionFeature_ARRAYSIZE; ++i) {
     proto::ModelExecutionFeature feature = proto::ModelExecutionFeature(i);
     if (!ShouldCheckSettingForFeature(feature)) {
       continue;
     }
 
-    bool is_enabled =
-        GetPrefState(feature) == prefs::FeatureOptInState::kEnabled;
     base::UmaHistogramBoolean(
         base::StrCat(
             {"OptimizationGuide.ModelExecution.FeatureEnabledAtStartup.",
              GetStringNameForModelExecutionFeature(feature)}),
-        is_enabled);
-    if (is_enabled) {
-      features_enabled_at_startup_.insert(static_cast<int>(feature));
-    }
+        ShouldFeatureBeCurrentlyEnabledForUser(feature));
   }
 }
 
@@ -520,11 +501,12 @@ void ModelExecutionFeaturesController::InitializePrefListener() {
         base::BindRepeating(
             &ModelExecutionFeaturesController::OnFeatureSettingPrefChanged,
             base::Unretained(this), feature));
+    pref_change_registrar_.Add(
+        model_execution::prefs::GetEnterprisePolicyPrefName(feature),
+        base::BindRepeating(&ModelExecutionFeaturesController::
+                                OnFeatureEnterprisePolicyPrefChanged,
+                            base::Unretained(this), feature));
   }
 }
 
-void ModelExecutionFeaturesController::SimulateBrowserRestartForTesting() {
-  InitializeFeatureSettings();
-}
-
 }  // namespace optimization_guide

components/optimization_guide/core/model_execution/model_execution_features_controller.h

@@ -80,8 +80,6 @@ class ModelExecutionFeaturesController
   // Removes `observer`.
   void RemoveObserver(SettingsEnabledObserver* observer);
 
-  void SimulateBrowserRestartForTesting();
-
  private:
   // Enumerates the reasons an user is invalid.
   enum class UserValidityResult {
@@ -97,6 +95,10 @@ class ModelExecutionFeaturesController
   // Called when the feature-specific toggle pref is changed.
   void OnFeatureSettingPrefChanged(proto::ModelExecutionFeature feature);
 
+  // Called when the feature-specific enterprise policy pref is changed.
+  void OnFeatureEnterprisePolicyPrefChanged(
+      proto::ModelExecutionFeature feature);
+
   void StartObservingAccountChanges();
 
   // signin::IdentityManager::Observer implementation:
@@ -129,10 +131,6 @@ class ModelExecutionFeaturesController
   // the conditions disallow the features.
   void ResetInvalidFeaturePrefs();
 
-  // Computed at the time `this` is constructed. Stores the set of features
-  // that were enabled at the time when browser started.
-  std::unordered_set<int> features_enabled_at_startup_;
-
   base::ScopedObservation<signin::IdentityManager,
                           ModelExecutionFeaturesController>
       identity_manager_observation_{this};

components/optimization_guide/core/model_execution/settings_enabled_observer.h

@@ -19,10 +19,18 @@ class SettingsEnabledObserver : public base::CheckedObserver {
 
   // Notifies `this` that the consumer feature team should prepare to enable
   // their feature when browser restarts. After browser restart, the feature
-  // team should call `ShouldFeatureBeCurrentlyEnabledForUser` before
-  // displaying any feature functionality.
+  // team should call `ShouldFeatureBeCurrentlyEnabledForUser` before displaying
+  // any feature functionality. TODO(rajendrant): Remove this once all the
+  // consumers stop using it.
   virtual void PrepareToEnableOnRestart() = 0;
 
+  // Notifies the consumers whenever the feature enabled state is changed.
+  // `is_now_enabled` indicates the current enabled state of the feature. This
+  // could be invoked without the enabled change toggling. This is not called
+  // automatically when the observer is added initially. Consumers should call
+  // `ShouldFeatureBeCurrentlyEnabledForUser` to check the feature state.
+  virtual void OnChangeInFeatureCurrentlyEnabledState(bool is_now_enabled) {}
+
   SettingsEnabledObserver(const SettingsEnabledObserver&) = delete;
   SettingsEnabledObserver& operator=(const SettingsEnabledObserver&) = delete;