add feature flag

Signed-off-by: Jiaping Zeng <[email protected]>

Author: jiapingzeng

common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java

@@ -483,4 +483,10 @@ private MLCommonsSettings() {}
     // Feature flag for streaming feature
     public static final Setting<Boolean> ML_COMMONS_STREAM_ENABLED = Setting
         .boolSetting(ML_PLUGIN_SETTING_PREFIX + "stream_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
+
+    // Feature flag for AG-UI agent support
+    public static final Setting<Boolean> ML_COMMONS_AG_UI_ENABLED = Setting
+        .boolSetting(ML_PLUGIN_SETTING_PREFIX + "ag_ui_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
+    public static final String ML_COMMONS_AG_UI_DISABLED_MESSAGE =
+        "The AG-UI agent feature is not enabled. To enable, please update the setting " + ML_COMMONS_AG_UI_ENABLED.getKey();
 }

common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java

@@ -22,6 +22,7 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_REMOTE_INFERENCE_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STREAM_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_ENABLED;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -63,6 +64,8 @@ public class MLFeatureEnabledSetting {
 
     private volatile Boolean isStreamEnabled;
 
+    private volatile Boolean isAGUIEnabled;
+
     private final List<SettingsChangeListener> listeners = new ArrayList<>();
 
     public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings) {
@@ -83,6 +86,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         isAgenticMemoryEnabled = ML_COMMONS_AGENTIC_MEMORY_ENABLED.get(settings);
         isIndexInsightEnabled = ML_COMMONS_INDEX_INSIGHT_FEATURE_ENABLED.get(settings);
         isStreamEnabled = ML_COMMONS_STREAM_ENABLED.get(settings);
+        isAGUIEnabled = ML_COMMONS_AG_UI_ENABLED.get(settings);
 
         clusterService
             .getClusterSettings()
@@ -112,6 +116,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         clusterService
             .getClusterSettings()
             .addSettingsUpdateConsumer(ML_COMMONS_INDEX_INSIGHT_FEATURE_ENABLED, it -> isIndexInsightEnabled = it);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_AG_UI_ENABLED, it -> isAGUIEnabled = it);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED, it -> {
             isStaticMetricCollectionEnabled = it;
             for (SettingsChangeListener listener : listeners) {
@@ -245,4 +250,12 @@ public boolean isIndexInsightEnabled() {
     public boolean isStreamEnabled() {
         return isStreamEnabled;
     }
+
+    /**
+     * Whether the AG-UI agent feature is enabled. If disabled, AG-UI agents will be blocked.
+     * @return whether the AG-UI agent feature is enabled.
+     */
+    public boolean isAGUIEnabled() {
+        return isAGUIEnabled;
+    }
 }

plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

@@ -1364,7 +1364,8 @@ public List<Setting<?>> getSettings() {
                 MLCommonsSettings.ML_COMMONS_INDEX_INSIGHT_FEATURE_ENABLED,
                 MLCommonsSettings.REMOTE_METADATA_GLOBAL_TENANT_ID,
                 MLCommonsSettings.REMOTE_METADATA_GLOBAL_RESOURCE_CACHE_TTL,
-                MLCommonsSettings.ML_COMMONS_STREAM_ENABLED
+                MLCommonsSettings.ML_COMMONS_STREAM_ENABLED,
+                MLCommonsSettings.ML_COMMONS_AG_UI_ENABLED
             );
         return settings;
     }

plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteAction.java

@@ -11,6 +11,7 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_EXECUTE_TOOL_DISABLED_MESSAGE;
 import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
 import static org.opensearch.ml.utils.MLExceptionUtils.AGENT_FRAMEWORK_DISABLED_ERR_MSG;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_AGENT_ID;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_ALGORITHM;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_TOOL_NAME;
@@ -128,6 +129,9 @@ MLExecuteTaskRequest getRequest(RestRequest request) throws IOException {
 
             String requestBodyJson = request.contentOrSourceParam().v2().utf8ToString();
             if (AGUIInputConverter.isAGUIInput(requestBodyJson)) {
+                if (!mlFeatureEnabledSetting.isAGUIEnabled()) {
+                    throw new IllegalStateException(ML_COMMONS_AG_UI_DISABLED_MESSAGE);
+                }
                 throw new IllegalArgumentException(
                     "AG-UI agents require streaming execution. "
                         + "Please use the streaming endpoint: POST /_plugins/_ml/agents/"

plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteStreamAction.java

@@ -16,6 +16,7 @@
 import static org.opensearch.ml.plugin.MachineLearningPlugin.STREAM_EXECUTE_THREAD_POOL;
 import static org.opensearch.ml.utils.MLExceptionUtils.AGENT_FRAMEWORK_DISABLED_ERR_MSG;
 import static org.opensearch.ml.utils.MLExceptionUtils.STREAM_DISABLED_ERR_MSG;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_AGENT_ID;
 import static org.opensearch.ml.utils.RestActionUtils.isAsync;
 import static org.opensearch.ml.utils.TenantAwareHelper.getTenantID;
@@ -361,6 +362,9 @@ MLExecuteTaskRequest getRequest(String agentId, RestRequest request, BytesRefere
         String requestBodyJson = content.utf8ToString();
         Input input;
         if (AGUIInputConverter.isAGUIInput(requestBodyJson)) {
+            if (!mlFeatureEnabledSetting.isAGUIEnabled()) {
+                throw new IllegalStateException(ML_COMMONS_AG_UI_DISABLED_MESSAGE);
+            }
             log.debug("AG-UI: Detected AG-UI input format for streaming agent: {}", agentId);
             input = AGUIInputConverter.convertFromAGUIInput(requestBodyJson, agentId, tenantId, async);
         } else {