Add feature flag to delete mem container (opensearch-project#4072)

Signed-off-by: Nathalie Jonathan <[email protected]>

Author: nathaliellenaa

plugin/src/main/java/org/opensearch/ml/action/memorycontainer/TransportDeleteMemoryContainerAction.java

@@ -6,6 +6,7 @@
 package org.opensearch.ml.action.memorycontainer;
 
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_CONTAINER_INDEX;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
 
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.ActionRequest;
@@ -73,6 +74,11 @@ public TransportDeleteMemoryContainerAction(
 
     @Override
     protected void doExecute(Task task, ActionRequest request, ActionListener<DeleteResponse> actionListener) {
+        if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
+            actionListener.onFailure(new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN));
+            return;
+        }
+
         MLMemoryContainerDeleteRequest deleteRequest = MLMemoryContainerDeleteRequest.fromActionRequest(request);
         String memoryContainerId = deleteRequest.getMemoryContainerId();
         String tenantId = deleteRequest.getTenantId();

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

@@ -7,12 +7,15 @@
 
 import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.BASE_MEMORY_CONTAINERS_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.TenantAwareHelper.getTenantID;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Locale;
 
+import org.opensearch.OpenSearchStatusException;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.memorycontainer.MLMemoryContainerDeleteAction;
 import org.opensearch.ml.common.transport.memorycontainer.MLMemoryContainerDeleteRequest;
@@ -53,6 +56,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);
+        }
+
         String memoryContainerId = request.param(PARAMETER_MEMORY_CONTAINER_ID);
         String tenantId = getTenantID(mlFeatureEnabledSetting.isMultiTenancyEnabled(), request);
         MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequest = new MLMemoryContainerDeleteRequest(memoryContainerId, tenantId);

plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportDeleteMemoryContainerActionTests.java

@@ -22,6 +22,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.opensearch.OpenSearchStatusException;
@@ -97,6 +98,9 @@ public class TransportDeleteMemoryContainerActionTests extends OpenSearchTestCas
     @Mock
     private MemoryContainerHelper memoryContainerHelper;
 
+    @Captor
+    private ArgumentCaptor<Exception> exceptionCaptor;
+
     TransportDeleteMemoryContainerAction transportDeleteMemoryContainerAction;
     MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequest;
     ThreadContext threadContext;
@@ -145,6 +149,7 @@ 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);
@@ -297,4 +302,21 @@ public void testDeleteMemoryContainer_UnauthorizedUser() {
                 )
             );
     }
+
+    public void testDoExecuteWithAgenticMemoryDisabled() throws InterruptedException {
+        // Disable agentic memory feature
+        when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(false);
+
+        // Execute
+        transportDeleteMemoryContainerAction.doExecute(null, mlMemoryContainerDeleteRequest, actionListener);
+
+        // Verify failure response due to feature being disabled
+        verify(actionListener).onFailure(exceptionCaptor.capture());
+        Exception exception = exceptionCaptor.getValue();
+        assertNotNull(exception);
+        assertTrue(exception instanceof OpenSearchStatusException);
+        OpenSearchStatusException statusException = (OpenSearchStatusException) exception;
+        assertEquals(RestStatus.FORBIDDEN, statusException.status());
+        assertEquals("The Agentic Memory APIs are not enabled. To enable, please update the setting plugins.ml_commons.agentic_memory_enabled", exception.getMessage());
+    }
 }

plugin/src/test/java/org/opensearch/ml/rest/RestMLDeleteMemoryContainerActionTests.java

@@ -21,6 +21,7 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.delete.DeleteResponse;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.action.ActionListener;
@@ -58,6 +59,7 @@ public class RestMLDeleteMemoryContainerActionTests extends OpenSearchTestCase {
     public void setup() {
         MockitoAnnotations.openMocks(this);
         when(mlFeatureEnabledSetting.isRemoteInferenceEnabled()).thenReturn(false);
+        when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true);
         restMLDeleteMemoryContainerAction = new RestMLDeleteMemoryContainerAction(mlFeatureEnabledSetting);
 
         threadPool = new TestThreadPool(this.getClass().getSimpleName() + "ThreadPool");
@@ -145,6 +147,19 @@ public void testPrepareRequest_MultiTenancyDisabled() throws Exception {
         assertNull(mlMemoryContainerDeleteRequest.getTenantId());
     }
 
+    public void testPrepareRequestWithAgenticMemoryDisabled() throws Exception {
+        // Disable agentic memory feature
+        when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(false);
+
+        RestRequest request = getRestRequest("memory_container_id", "tenant_id");
+
+        // Expect OpenSearchStatusException when feature is disabled
+        thrown.expect(OpenSearchStatusException.class);
+        thrown.expectMessage("The Agentic Memory APIs are not enabled. To enable, please update the setting plugins.ml_commons.agentic_memory_enabled");
+
+        restMLDeleteMemoryContainerAction.handleRequest(request, channel, client);
+    }
+
     public void testActionNameConstant() {
         // Test that the action name constant is correctly defined
         assertEquals("ml_delete_memory_container_action", restMLDeleteMemoryContainerAction.getName());