Add protos for text safety feature

Bug: b:330346344
Change-Id: I63426ba9fe993bb5e6530f8dee0e92f95545e16a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5380976
Reviewed-by: Robert Ogden <robertogden@chromium.org>
Commit-Queue: Sophie Chang <sophiechang@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1275168}
diff --git a/components/optimization_guide/core/model_execution/model_execution_features.cc b/components/optimization_guide/core/model_execution/model_execution_features.cc
index 097e7b7..39622d8f 100644
--- a/components/optimization_guide/core/model_execution/model_execution_features.cc
+++ b/components/optimization_guide/core/model_execution/model_execution_features.cc
@@ -57,6 +57,7 @@
       is_graduated = base::FeatureList::IsEnabled(kWallpaperSearchGraduated);
       break;
     // Non-features.
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
       NOTREACHED();
@@ -79,6 +80,7 @@
       return &kTabOrganizationSettingsVisibility;
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH:
       return &kWallpaperSearchSettingsVisibility;
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
       NOTREACHED();
@@ -101,6 +103,10 @@
         proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST) {
       continue;
     }
+    if (model_execution_feature ==
+        proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY) {
+      continue;
+    }
     const auto* feature =
         GetFeatureToUseToCheckSettingsVisibility(model_execution_feature);
     if (GetFieldTrialParamByFeatureAsBool(*feature, "allow_unsigned_user",
diff --git a/components/optimization_guide/core/model_execution/model_execution_features_controller.cc b/components/optimization_guide/core/model_execution/model_execution_features_controller.cc
index 13d3917..e184d20 100644
--- a/components/optimization_guide/core/model_execution/model_execution_features_controller.cc
+++ b/components/optimization_guide/core/model_execution/model_execution_features_controller.cc
@@ -21,7 +21,8 @@
 
 bool ShouldCheckSettingForFeature(proto::ModelExecutionFeature feature) {
   return feature != proto::MODEL_EXECUTION_FEATURE_UNSPECIFIED &&
-         feature != proto::MODEL_EXECUTION_FEATURE_TEST;
+         feature != proto::MODEL_EXECUTION_FEATURE_TEST &&
+         feature != proto::MODEL_EXECUTION_FEATURE_TEXT_SAFETY;
 }
 
 // Util class for recording the construction and validation of Settings
diff --git a/components/optimization_guide/core/model_execution/model_execution_fetcher.cc b/components/optimization_guide/core/model_execution/model_execution_fetcher.cc
index 96564a8..634ad41 100644
--- a/components/optimization_guide/core/model_execution/model_execution_fetcher.cc
+++ b/components/optimization_guide/core/model_execution/model_execution_fetcher.cc
@@ -131,6 +131,9 @@
             }
           }
         })");
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
+      // TODO: b/330346344 - Add traffic annotation.
+      return MISSING_TRAFFIC_ANNOTATION;
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
       // Used for testing purposes. No real features use this.
       return MISSING_TRAFFIC_ANNOTATION;
diff --git a/components/optimization_guide/core/model_execution/model_execution_prefs.cc b/components/optimization_guide/core/model_execution/model_execution_prefs.cc
index 9170025d..3f71f00 100644
--- a/components/optimization_guide/core/model_execution/model_execution_prefs.cc
+++ b/components/optimization_guide/core/model_execution/model_execution_prefs.cc
@@ -43,6 +43,7 @@
       return kTabOrganizationEnterprisePolicyAllowed;
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH:
       return kWallpaperSearchEnterprisePolicyAllowed;
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
       NOTREACHED();
diff --git a/components/optimization_guide/core/model_execution/model_execution_util.cc b/components/optimization_guide/core/model_execution/model_execution_util.cc
index be5a80e..7896a45 100644
--- a/components/optimization_guide/core/model_execution/model_execution_util.cc
+++ b/components/optimization_guide/core/model_execution/model_execution_util.cc
@@ -27,8 +27,9 @@
       SetExecutionRequestTemplate<ComposeFeatureTypeMap>(log_ai_request,
                                                          request_metadata);
       return;
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
-      // Do not log request for test.
+      // Do not log request for test and text safety.
       return;
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
       // Don't log any request data when the feature is not specified.
@@ -55,8 +56,9 @@
       SetExecutionResponseTemplate<ComposeFeatureTypeMap>(log_ai_request,
                                                           response_metadata);
       return;
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
-      // Do not log response for test.
+      // Do not log response for test and text safety.
       return;
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
       // Don't log any response data when the feature is not specified.
diff --git a/components/optimization_guide/core/optimization_guide_features.cc b/components/optimization_guide/core/optimization_guide_features.cc
index 8585138e..df7274c 100644
--- a/components/optimization_guide/core/optimization_guide_features.cc
+++ b/components/optimization_guide/core/optimization_guide_features.cc
@@ -299,6 +299,7 @@
       break;
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
       // Logging disabled.
       NOTREACHED();
       break;
diff --git a/components/optimization_guide/core/optimization_guide_prefs.cc b/components/optimization_guide/core/optimization_guide_prefs.cc
index da047f8d..7701dbd7 100644
--- a/components/optimization_guide/core/optimization_guide_prefs.cc
+++ b/components/optimization_guide/core/optimization_guide_prefs.cc
@@ -76,6 +76,7 @@
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_WALLPAPER_SEARCH:
       return "optimization_guide.wallpaper_search_setting_state";
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
       NOTREACHED();
       return "Invalid";
@@ -93,6 +94,7 @@
         static_cast<proto::ModelExecutionFeature>(i);
     switch (feature) {
       case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
+      case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
       case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
         continue;
       default:
diff --git a/components/optimization_guide/core/optimization_guide_util.cc b/components/optimization_guide/core/optimization_guide_util.cc
index 3ff42582..4319a29 100644
--- a/components/optimization_guide/core/optimization_guide_util.cc
+++ b/components/optimization_guide/core/optimization_guide_util.cc
@@ -88,6 +88,8 @@
       return "Compose";
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEST:
       return "Test";
+    case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_TEXT_SAFETY:
+      return "TextSafety";
     case proto::ModelExecutionFeature::MODEL_EXECUTION_FEATURE_UNSPECIFIED:
       return "Unknown";
       // Must be in sync with the ModelExecutionFeature variant in
diff --git a/components/optimization_guide/proto/model_execution.proto b/components/optimization_guide/proto/model_execution.proto
index 68cbab7..3e19b2f1 100644
--- a/components/optimization_guide/proto/model_execution.proto
+++ b/components/optimization_guide/proto/model_execution.proto
@@ -61,6 +61,10 @@
   // prototyping. Please reach out to optimization_guide OWNERS when looking to
   // use this enum or want to go beyond prototyping.
   MODEL_EXECUTION_FEATURE_TEST = 4;
+  // This feature does not map to a user-visible feature that is being
+  // launched, but instead is an internal detail to support other model
+  // execution features.
+  MODEL_EXECUTION_FEATURE_TEXT_SAFETY = 5;
 }
 
 message OnDeviceModelExecutionConfig {
diff --git a/components/optimization_guide/proto/model_quality_metadata.proto b/components/optimization_guide/proto/model_quality_metadata.proto
index e4fddf5c..8dae068 100644
--- a/components/optimization_guide/proto/model_quality_metadata.proto
+++ b/components/optimization_guide/proto/model_quality_metadata.proto
@@ -166,6 +166,9 @@
 message TextSafetyModelRequest {
   // The text sent to the text safety model for evaluation.
   string text = 1;
+
+  // The URL sent to the text safety model for evaluation, if any.
+  string url = 2;
 }
 
 message TextSafetyModelResponse {
@@ -174,4 +177,8 @@
 
   // Whether the output was deemed unsafe.
   bool is_unsafe = 2;
+
+  // The ID for the execution of the text safety model if it leveraged the
+  // server for text safety evaluation.
+  string server_execution_id = 3;
 }
diff --git a/tools/metrics/histograms/metadata/optimization/histograms.xml b/tools/metrics/histograms/metadata/optimization/histograms.xml
index aa87913..11fdaae 100644
--- a/tools/metrics/histograms/metadata/optimization/histograms.xml
+++ b/tools/metrics/histograms/metadata/optimization/histograms.xml
@@ -32,6 +32,7 @@
   <variant name="Compose" summary="Compose"/>
   <variant name="TabOrganization" summary="Tab Organization"/>
   <variant name="Test" summary="Test"/>
+  <variant name="TextSafety" summary="Text Safety"/>
   <variant name="Unknown" summary="Unknown"/>
   <variant name="WallpaperSearch" summary="Wallpaper search"/>
 </variants>