adding feature settings for all the agentic memory apis (#4074)

dhrubo-os · web-flow · commit 7ade595d415e · 2025-08-06T15:58:46.000-07:00
Signed-off-by: Dhrubo Saha &lt;dhrubo@amazon.com&gt;
diff --git a/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportAddMemoriesAction.java b/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportAddMemoriesAction.java
@@ -14,6 +14,7 @@
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_EMBEDDING_FIELD;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_FIELD;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PERSONAL_INFORMATION_ORGANIZER_PROMPT;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
 
 import java.time.Instant;
 import java.util.ArrayList;
@@ -173,6 +174,11 @@ public TransportAddMemoriesAction(
 
     @Override
     protected void doExecute(Task task, MLAddMemoriesRequest request, ActionListener<MLAddMemoriesResponse> actionListener) {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            actionListener.onFailure(new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN));
+            return;
+        }
+
         User user = RestActionUtils.getUserContext(client);
         MLAddMemoriesInput input = request.getMlAddMemoryInput();
 
diff --git a/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportDeleteMemoryAction.java b/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportDeleteMemoryAction.java
@@ -17,6 +17,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
+import org.opensearch.ml.common.settings.MLCommonsSettings;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLDeleteMemoryAction;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLDeleteMemoryRequest;
@@ -65,6 +66,14 @@ public TransportDeleteMemoryAction(
 
     @Override
     protected void doExecute(Task task, ActionRequest request, ActionListener<DeleteResponse> actionListener) {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            actionListener
+                .onFailure(
+                    new OpenSearchStatusException(MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN)
+                );
+            return;
+        }
+
         MLDeleteMemoryRequest deleteRequest = MLDeleteMemoryRequest.fromActionRequest(request);
         String memoryContainerId = deleteRequest.getMemoryContainerId();
         String memoryId = deleteRequest.getMemoryId();
diff --git a/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportSearchMemoriesAction.java b/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportSearchMemoriesAction.java
@@ -6,6 +6,7 @@
 package org.opensearch.ml.action.memorycontainer.memory;
 
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.*;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
 
 import java.io.IOException;
 import java.time.Instant;
@@ -78,6 +79,11 @@ public TransportSearchMemoriesAction(
 
     @Override
     protected void doExecute(Task task, MLSearchMemoriesRequest request, ActionListener<MLSearchMemoriesResponse> actionListener) {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            actionListener.onFailure(new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN));
+            return;
+        }
+
         MLSearchMemoriesInput input = request.getMlSearchMemoriesInput();
         String tenantId = request.getTenantId();
 
diff --git a/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportUpdateMemoryAction.java b/plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportUpdateMemoryAction.java
@@ -25,6 +25,7 @@
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.memorycontainer.MemoryStorageConfig;
+import org.opensearch.ml.common.settings.MLCommonsSettings;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLUpdateMemoryAction;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLUpdateMemoryRequest;
@@ -81,6 +82,14 @@ public TransportUpdateMemoryAction(
 
     @Override
     protected void doExecute(Task task, ActionRequest request, ActionListener<UpdateResponse> actionListener) {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            actionListener
+                .onFailure(
+                    new OpenSearchStatusException(MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN)
+                );
+            return;
+        }
+
         MLUpdateMemoryRequest updateRequest = MLUpdateMemoryRequest.fromActionRequest(request);
         String memoryContainerId = updateRequest.getMemoryContainerId();
         String memoryId = updateRequest.getMemoryId();
diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -934,13 +934,13 @@ public List<RestHandler> getRestHandlers(
             mlFeatureEnabledSetting
         );
         RestMLGetMemoryContainerAction restMLGetMemoryContainerAction = new RestMLGetMemoryContainerAction(mlFeatureEnabledSetting);
+        RestMLAddMemoriesAction restMLAddMemoriesAction = new RestMLAddMemoriesAction(mlFeatureEnabledSetting);
+        RestMLSearchMemoriesAction restMLSearchMemoriesAction = new RestMLSearchMemoriesAction(mlFeatureEnabledSetting);
+        RestMLDeleteMemoryAction restMLDeleteMemoryAction = new RestMLDeleteMemoryAction(mlFeatureEnabledSetting);
+        RestMLUpdateMemoryAction restMLUpdateMemoryAction = new RestMLUpdateMemoryAction(mlFeatureEnabledSetting);
         RestMLDeleteMemoryContainerAction restMLDeleteMemoryContainerAction = new RestMLDeleteMemoryContainerAction(
             mlFeatureEnabledSetting
         );
-        RestMLAddMemoriesAction restMLAddMemoriesAction = new RestMLAddMemoriesAction();
-        RestMLSearchMemoriesAction restMLSearchMemoriesAction = new RestMLSearchMemoriesAction(mlFeatureEnabledSetting);
-        RestMLDeleteMemoryAction restMLDeleteMemoryAction = new RestMLDeleteMemoryAction();
-        RestMLUpdateMemoryAction restMLUpdateMemoryAction = new RestMLUpdateMemoryAction();
         RestMemorySearchConversationsAction restSearchConversationsAction = new RestMemorySearchConversationsAction(
             mlFeatureEnabledSetting
         );
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLAddMemoriesAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLAddMemoriesAction.java
@@ -7,12 +7,16 @@
 
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORIES_PATH;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETER_MEMORY_CONTAINER_ID;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
 
 import java.io.IOException;
 import java.util.List;
 
+import org.opensearch.OpenSearchStatusException;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLAddMemoriesAction;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLAddMemoriesInput;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLAddMemoriesRequest;
@@ -29,6 +33,14 @@
 public class RestMLAddMemoriesAction extends BaseRestHandler {
 
     private static final String ML_ADD_MEMORIES_ACTION = "ml_add_memories_action";
+    private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
+
+    /**
+     * Constructor
+     */
+    public RestMLAddMemoriesAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
+    }
 
     @Override
     public List<Route> routes() {
@@ -42,6 +54,10 @@ public String getName() {
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            throw new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN);
+        }
+
         MLAddMemoriesRequest mlAddMemoryRequest = getRequest(request);
         return channel -> client.execute(MLAddMemoriesAction.INSTANCE, mlAddMemoryRequest, new RestToXContentListener<>(channel));
     }
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLDeleteMemoryAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLDeleteMemoryAction.java
@@ -8,11 +8,15 @@
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.DELETE_MEMORY_PATH;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETER_MEMORY_CONTAINER_ID;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETER_MEMORY_ID;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
 
 import java.io.IOException;
 import java.util.List;
 
+import org.opensearch.OpenSearchStatusException;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLDeleteMemoryAction;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLDeleteMemoryRequest;
 import org.opensearch.rest.BaseRestHandler;
@@ -29,11 +33,14 @@
 public class RestMLDeleteMemoryAction extends BaseRestHandler {
 
     private static final String ML_DELETE_MEMORY_ACTION = "ml_delete_memory_action";
+    private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     /**
      * Constructor
      */
-    public RestMLDeleteMemoryAction() {}
+    public RestMLDeleteMemoryAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
+    }
 
     @Override
     public String getName() {
@@ -47,6 +54,10 @@ public List<Route> routes() {
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            throw new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN);
+        }
+
         MLDeleteMemoryRequest mlDeleteMemoryRequest = getRequest(request);
         return channel -> client.execute(MLDeleteMemoryAction.INSTANCE, mlDeleteMemoryRequest, new RestToXContentListener<>(channel));
     }
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLSearchMemoriesAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLSearchMemoriesAction.java
@@ -8,12 +8,15 @@
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETER_MEMORY_CONTAINER_ID;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.SEARCH_MEMORIES_PATH;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
 import static org.opensearch.ml.utils.TenantAwareHelper.getTenantID;
 
 import java.io.IOException;
 import java.util.List;
 
+import org.opensearch.OpenSearchStatusException;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLSearchMemoriesAction;
@@ -55,6 +58,10 @@ public List<Route> routes() {
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            throw new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN);
+        }
+
         MLSearchMemoriesRequest mlSearchMemoriesRequest = getRequest(request);
         return channel -> client.execute(MLSearchMemoriesAction.INSTANCE, mlSearchMemoriesRequest, new RestToXContentListener<>(channel));
     }
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLUpdateMemoryAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLUpdateMemoryAction.java
@@ -9,12 +9,16 @@
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETER_MEMORY_CONTAINER_ID;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETER_MEMORY_ID;
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.UPDATE_MEMORY_PATH;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
 
 import java.io.IOException;
 import java.util.List;
 
+import org.opensearch.OpenSearchStatusException;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLUpdateMemoryAction;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLUpdateMemoryInput;
 import org.opensearch.ml.common.transport.memorycontainer.memory.MLUpdateMemoryRequest;
@@ -32,11 +36,14 @@
 public class RestMLUpdateMemoryAction extends BaseRestHandler {
 
     private static final String ML_UPDATE_MEMORY_ACTION = "ml_update_memory_action";
+    private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     /**
      * Constructor
      */
-    public RestMLUpdateMemoryAction() {}
+    public RestMLUpdateMemoryAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
+    }
 
     @Override
     public String getName() {
@@ -50,6 +57,10 @@ public List<Route> routes() {
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            throw new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN);
+        }
+
         MLUpdateMemoryRequest mlUpdateMemoryRequest = getRequest(request);
         return channel -> client.execute(MLUpdateMemoryAction.INSTANCE, mlUpdateMemoryRequest, new RestToXContentListener<>(channel));
     }
diff --git a/plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportCreateMemoryContainerActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportCreateMemoryContainerActionTests.java
@@ -1087,14 +1087,14 @@ public void testDoExecuteWithNullMemoryStorageConfig() throws InterruptedExcepti
     public void testDoExecuteWithTenantValidationFailure() throws InterruptedException {
         // Enable multi-tenancy and provide null tenant ID to trigger validation failure
         when(mlFeatureEnabledSetting.isMultiTenancyEnabled()).thenReturn(true);
-        
+
         MLCreateMemoryContainerInput invalidTenantInput = MLCreateMemoryContainerInput
-            .builder()
-            .name("invalid-tenant-container")
-            .description("Container with invalid tenant")
-            .memoryStorageConfig(memoryStorageConfig)
-            .tenantId(null) // This should trigger tenant validation failure
-            .build();
+                .builder()
+                .name("invalid-tenant-container")
+                .description("Container with invalid tenant")
+                .memoryStorageConfig(memoryStorageConfig)
+                .tenantId(null) // This should trigger tenant validation failure
+                .build();
 
         MLCreateMemoryContainerRequest invalidTenantRequest = new MLCreateMemoryContainerRequest(invalidTenantInput);
 
diff --git a/plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportDeleteMemoryContainerActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportDeleteMemoryContainerActionTests.java
@@ -149,11 +149,13 @@ public void setup() throws IOException {
             listener.onResponse(mockContainer);
             return null;
         }).when(memoryContainerHelper).getMemoryContainer(any(), any());
-        when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true);
 
         threadContext = new ThreadContext(settings);
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
+
+        // Setup ML feature settings
+        when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); // Enable by default for tests
     }
 
     @Test
@@ -307,10 +309,14 @@ public void testDoExecuteWithAgenticMemoryDisabled() throws InterruptedException
         // Disable agentic memory feature
         when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(false);
 
+        // Create request
+        MLMemoryContainerDeleteRequest request = MLMemoryContainerDeleteRequest.builder().memoryContainerId(MEMORY_CONTAINER_ID).build();
+
         // Execute
-        transportDeleteMemoryContainerAction.doExecute(null, mlMemoryContainerDeleteRequest, actionListener);
+        transportDeleteMemoryContainerAction.doExecute(null, request, actionListener);
 
         // Verify failure response due to feature being disabled
+        ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
         verify(actionListener).onFailure(exceptionCaptor.capture());
         Exception exception = exceptionCaptor.getValue();
         assertNotNull(exception);
diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLCreateMemoryContainerActionTests.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLCreateMemoryContainerActionTests.java
@@ -152,14 +152,14 @@ public void testGetRequestWithMinimalInput() throws IOException {
 
     public void testGetRequestWithMultiTenancyEnabled() throws IOException {
         when(mlFeatureEnabledSetting.isMultiTenancyEnabled()).thenReturn(true);
-        
+
         String requestBody = "{\n" +
-            "  \"name\": \"tenant-container\"\n" +
-            "}";
+                "  \"name\": \"tenant-container\"\n" +
+                "}";
 
         Map<String, String> headers = new HashMap<>();
         headers.put(Constants.TENANT_ID_HEADER, "test-tenant");
-        
+
         RestRequest request = createRestRequestWithHeaders(requestBody, headers);
         MLCreateMemoryContainerRequest mlCreateMemoryContainerRequest = restMLCreateMemoryContainerAction.getRequest(request);
 
@@ -301,14 +301,14 @@ public void testGetRequestWithUnknownFields() throws IOException {
 
     public void testMultiTenancyDisabledTenantIdIsNull() throws IOException {
         when(mlFeatureEnabledSetting.isMultiTenancyEnabled()).thenReturn(false);
-        
+
         String requestBody = "{\n" +
-            "  \"name\": \"no-tenant-container\"\n" +
-            "}";
+                "  \"name\": \"no-tenant-container\"\n" +
+                "}";
 
         Map<String, String> headers = new HashMap<>();
         headers.put(Constants.TENANT_ID_HEADER, "should-be-ignored");
-        
+
         RestRequest request = createRestRequestWithHeaders(requestBody, headers);
         MLCreateMemoryContainerRequest mlCreateMemoryContainerRequest = restMLCreateMemoryContainerAction.getRequest(request);
 
@@ -369,9 +369,9 @@ public void testPrepareRequestWithAgenticMemoryDisabled() throws IOException {
         when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(false);
 
         String requestContent = "{\n"
-            + "  \"name\": \"test-memory-container\",\n"
-            + "  \"description\": \"Test memory container description\"\n"
-            + "}";
+                + "  \"name\": \"test-memory-container\",\n"
+                + "  \"description\": \"Test memory container description\"\n"
+                + "}";
 
         RestRequest request = createRestRequest(requestContent);
 
diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLGetMemoryContainerActionTests.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLGetMemoryContainerActionTests.java
