Skip to content

Commit 02601dc

Browse files
dhrubo-osdhrubo-os
authored
adding feature flag for agentic memory (opensearch-project#4067)
Signed-off-by: Dhrubo Saha <[email protected]>
1 parent 4122b38 commit 02601dc

File tree

14 files changed

+188
-2
lines changed

14 files changed

+188
-2
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,4 +362,10 @@ private MLCommonsSettings() {}
362362
// Feature flag for enabling telemetry static metric collection job -- MLStatsJobProcessor
363363
public static final Setting<Boolean> ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED = Setting
364364
.boolSetting("plugins.ml_commons.metrics_static_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Final);
365+
366+
// Feature flag for Agentic memory APIs
367+
public static final Setting<Boolean> ML_COMMONS_AGENTIC_MEMORY_ENABLED = Setting
368+
.boolSetting("plugins.ml_commons.agentic_memory_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
369+
public static final String ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE =
370+
"The Agentic Memory APIs are not enabled. To enable, please update the setting " + ML_COMMONS_AGENTIC_MEMORY_ENABLED.getKey();
365371
}

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package org.opensearch.ml.common.settings;
77

8+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED;
89
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_ENABLED;
910
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENT_FRAMEWORK_ENABLED;
1011
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED;
@@ -58,6 +59,8 @@ public class MLFeatureEnabledSetting {
5859

5960
private volatile Boolean isMcpConnectorEnabled;
6061

62+
private volatile Boolean isAgenticMemoryEnabled;
63+
6164
private final List<SettingsChangeListener> listeners = new ArrayList<>();
6265

6366
public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings) {
@@ -76,6 +79,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
7679
isExecuteToolEnabled = ML_COMMONS_EXECUTE_TOOL_ENABLED.get(settings);
7780
isAgenticSearchEnabled = ML_COMMONS_AGENTIC_SEARCH_ENABLED.get(settings);
7881
isMcpConnectorEnabled = ML_COMMONS_MCP_CONNECTOR_ENABLED.get(settings);
82+
isAgenticMemoryEnabled = ML_COMMONS_AGENTIC_MEMORY_ENABLED.get(settings);
7983

8084
clusterService
8185
.getClusterSettings()
@@ -101,6 +105,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
101105
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_EXECUTE_TOOL_ENABLED, it -> isExecuteToolEnabled = it);
102106
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_AGENTIC_SEARCH_ENABLED, it -> isAgenticSearchEnabled = it);
103107
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MCP_CONNECTOR_ENABLED, it -> isMcpConnectorEnabled = it);
108+
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_AGENTIC_MEMORY_ENABLED, it -> isMcpConnectorEnabled = it);
104109
}
105110

106111
/**
@@ -199,6 +204,14 @@ public boolean isToolExecuteEnabled() {
199204
return isExecuteToolEnabled;
200205
}
201206

207+
/**
208+
* Whether the Agentic memory APIs are enabled. If disabled, Agentic memory APIs in ml-commons will be blocked
209+
* @return whether the agentic memory feature is enabled.
210+
*/
211+
public boolean isAgenticMemoryEnabled() {
212+
return isAgenticMemoryEnabled;
213+
}
214+
202215
@VisibleForTesting
203216
public void notifyMultiTenancyListeners(boolean isEnabled) {
204217
for (SettingsChangeListener listener : listeners) {

‎common/src/test/java/org/opensearch/ml/common/settings/MLCommonsSettingsTests.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,4 +70,36 @@ public void testRemoteInferenceEnabledByDefault() {
7070
public void testAllowModelUrlDisabledByDefault() {
7171
assertFalse(MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL.getDefault(null));
7272
}
73+
74+
@Test
75+
public void testAgenticMemoryDisabledByDefault() {
76+
assertFalse(MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED.getDefault(null));
77+
}
78+
79+
@Test
80+
public void testAgenticMemorySettingProperties() {
81+
// Test setting key
82+
assertEquals("plugins.ml_commons.agentic_memory_enabled", MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED.getKey());
83+
84+
// Test setting is dynamic
85+
assertTrue(
86+
MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED
87+
.getProperties()
88+
.contains(org.opensearch.common.settings.Setting.Property.Dynamic)
89+
);
90+
91+
// Test setting is node scope
92+
assertTrue(
93+
MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED
94+
.getProperties()
95+
.contains(org.opensearch.common.settings.Setting.Property.NodeScope)
96+
);
97+
}
98+
99+
@Test
100+
public void testAgenticMemoryDisabledMessage() {
101+
String expectedMessage =
102+
"The Agentic Memory APIs are not enabled. To enable, please update the setting plugins.ml_commons.agentic_memory_enabled";
103+
assertEquals(expectedMessage, MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE);
104+
}
73105
}

‎common/src/test/java/org/opensearch/ml/common/settings/MLFeatureEnabledSettingTests.java‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,8 @@ public void setUp() {
4646
MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED,
4747
MLCommonsSettings.ML_COMMONS_EXECUTE_TOOL_ENABLED,
4848
MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_ENABLED,
49-
MLCommonsSettings.ML_COMMONS_MCP_CONNECTOR_ENABLED
49+
MLCommonsSettings.ML_COMMONS_MCP_CONNECTOR_ENABLED,
50+
MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED
5051
)
5152
);
5253
when(mockClusterService.getClusterSettings()).thenReturn(mockClusterSettings);
@@ -70,6 +71,7 @@ public void testDefaults_allFeaturesEnabled() {
7071
.put("plugins.ml_commons.metrics_static_collection_enabled", true)
7172
.put("plugins.ml_commons.mcp_connector_enabled", true)
7273
.put("plugins.ml_commons.agentic_search_enabled", true)
74+
.put("plugins.ml_commons.agentic_memory_enabled", true)
7375
.build();
7476

7577
MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
@@ -88,6 +90,7 @@ public void testDefaults_allFeaturesEnabled() {
8890
assertTrue(setting.isStaticMetricCollectionEnabled());
8991
assertTrue(setting.isMcpConnectorEnabled());
9092
assertTrue(setting.isAgenticSearchEnabled());
93+
assertTrue(setting.isAgenticMemoryEnabled());
9194
}
9295

9396
@Test
@@ -108,6 +111,7 @@ public void testDefaults_someFeaturesDisabled() {
108111
.put("plugins.ml_commons.metrics_static_collection_enabled", false)
109112
.put("plugins.ml_commons.mcp_connector_enabled", false)
110113
.put("plugins.ml_commons.agentic_search_enabled", false)
114+
.put("plugins.ml_commons.agentic_memory_enabled", false)
111115
.build();
112116

113117
MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
@@ -126,6 +130,7 @@ public void testDefaults_someFeaturesDisabled() {
126130
assertFalse(setting.isStaticMetricCollectionEnabled());
127131
assertFalse(setting.isMcpConnectorEnabled());
128132
assertFalse(setting.isAgenticSearchEnabled());
133+
assertFalse(setting.isAgenticMemoryEnabled());
129134
}
130135

131136
@Test
@@ -140,4 +145,29 @@ public void testMultiTenancyChangeNotifiesListeners() {
140145
setting.notifyMultiTenancyListeners(true);
141146
verify(mockListener).onMultiTenancyEnabledChanged(true);
142147
}
148+
149+
@Test
150+
public void testAgenticMemoryEnabledByDefault() {
151+
Settings settings = Settings.EMPTY;
152+
MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
153+
154+
// Should be disabled by default
155+
assertFalse(setting.isAgenticMemoryEnabled());
156+
}
157+
158+
@Test
159+
public void testAgenticMemoryCanBeEnabled() {
160+
Settings settings = Settings.builder().put("plugins.ml_commons.agentic_memory_enabled", true).build();
161+
162+
MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
163+
assertTrue(setting.isAgenticMemoryEnabled());
164+
}
165+
166+
@Test
167+
public void testAgenticMemoryCanBeDisabled() {
168+
Settings settings = Settings.builder().put("plugins.ml_commons.agentic_memory_enabled", false).build();
169+
170+
MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
171+
assertFalse(setting.isAgenticMemoryEnabled());
172+
}
143173
}

‎plugin/src/main/java/org/opensearch/ml/action/memorycontainer/TransportCreateMemoryContainerAction.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,14 @@
2929
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.STATIC_MEMORY_INDEX_PREFIX;
3030
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.TAGS_FIELD;
3131
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.USER_ID_FIELD;
32+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
3233

3334
import java.time.Instant;
3435
import java.util.HashMap;
3536
import java.util.Locale;
3637
import java.util.Map;
3738

39+
import org.opensearch.OpenSearchStatusException;
3840
import org.opensearch.action.DocWriteResponse;
3941
import org.opensearch.action.index.IndexResponse;
4042
import org.opensearch.action.support.ActionFilters;
@@ -44,6 +46,7 @@
4446
import org.opensearch.common.util.concurrent.ThreadContext;
4547
import org.opensearch.commons.authuser.User;
4648
import org.opensearch.core.action.ActionListener;
49+
import org.opensearch.core.rest.RestStatus;
4750
import org.opensearch.ml.common.FunctionName;
4851
import org.opensearch.ml.common.MLIndex;
4952
import org.opensearch.ml.common.MLModel;
@@ -107,6 +110,11 @@ public TransportCreateMemoryContainerAction(
107110

108111
@Override
109112
protected void doExecute(Task task, MLCreateMemoryContainerRequest request, ActionListener<MLCreateMemoryContainerResponse> listener) {
113+
if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
114+
listener.onFailure(new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN));
115+
return;
116+
}
117+
110118
MLCreateMemoryContainerInput input = request.getMlCreateMemoryContainerInput();
111119

112120
// Validate tenant ID

‎plugin/src/main/java/org/opensearch/ml/action/memorycontainer/TransportGetMemoryContainerAction.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.common.xcontent.json.JsonXContent.jsonXContent;
99
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
1010
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_CONTAINER_INDEX;
11+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
1112

1213
import org.opensearch.ExceptionsHelper;
1314
import org.opensearch.OpenSearchStatusException;
@@ -82,6 +83,11 @@ public TransportGetMemoryContainerAction(
8283

8384
@Override
8485
protected void doExecute(Task task, ActionRequest request, ActionListener<MLMemoryContainerGetResponse> actionListener) {
86+
if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
87+
actionListener.onFailure(new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN));
88+
return;
89+
}
90+
8591
MLMemoryContainerGetRequest mlMemoryContainerGetRequest = MLMemoryContainerGetRequest.fromActionRequest(request);
8692
String memoryContainerId = mlMemoryContainerGetRequest.getMemoryContainerId();
8793
String tenantId = mlMemoryContainerGetRequest.getTenantId();

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1205,7 +1205,8 @@ public List<Setting<?>> getSettings() {
12051205
MLCommonsSettings.ML_COMMONS_METRIC_COLLECTION_ENABLED,
12061206
MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED,
12071207
MLCommonsSettings.ML_COMMONS_EXECUTE_TOOL_ENABLED,
1208-
MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_ENABLED
1208+
MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_ENABLED,
1209+
MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED
12091210
);
12101211
return settings;
12111212
}

‎plugin/src/main/java/org/opensearch/ml/rest/RestMLCreateMemoryContainerAction.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,15 @@
66
package org.opensearch.ml.rest;
77

88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
9+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
910
import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
1011

1112
import java.io.IOException;
1213
import java.util.List;
1314
import java.util.Locale;
1415

16+
import org.opensearch.OpenSearchStatusException;
17+
import org.opensearch.core.rest.RestStatus;
1518
import org.opensearch.core.xcontent.XContentParser;
1619
import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
1720
import org.opensearch.ml.common.transport.memorycontainer.MLCreateMemoryContainerAction;
@@ -54,6 +57,9 @@ public List<Route> routes() {
5457

5558
@Override
5659
public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
60+
if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
61+
throw new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN);
62+
}
5763
MLCreateMemoryContainerRequest mlCreateMemoryContainerRequest = getRequest(request);
5864
return channel -> client
5965
.execute(MLCreateMemoryContainerAction.INSTANCE, mlCreateMemoryContainerRequest, new RestToXContentListener<>(channel));

‎plugin/src/main/java/org/opensearch/ml/rest/RestMLGetMemoryContainerAction.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,16 @@
77

88
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.BASE_MEMORY_CONTAINERS_PATH;
99
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETER_MEMORY_CONTAINER_ID;
10+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE;
1011
import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
1112
import static org.opensearch.ml.utils.TenantAwareHelper.getTenantID;
1213

1314
import java.io.IOException;
1415
import java.util.List;
1516
import java.util.Locale;
1617

18+
import org.opensearch.OpenSearchStatusException;
19+
import org.opensearch.core.rest.RestStatus;
1720
import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
1821
import org.opensearch.ml.common.transport.memorycontainer.MLMemoryContainerGetAction;
1922
import org.opensearch.ml.common.transport.memorycontainer.MLMemoryContainerGetRequest;
@@ -54,6 +57,9 @@ public List<Route> routes() {
5457

5558
@Override
5659
public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
60+
if (!mlFeatureEnabledSetting.isAgenticMemoryEnabled()) {
61+
throw new OpenSearchStatusException(ML_COMMONS_AGENTIC_MEMORY_DISABLED_MESSAGE, RestStatus.FORBIDDEN);
62+
}
5763
MLMemoryContainerGetRequest mlMemoryContainerGetRequest = getRequest(request);
5864
return channel -> client
5965
.execute(MLMemoryContainerGetAction.INSTANCE, mlMemoryContainerGetRequest, new RestToXContentListener<>(channel));

‎plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportCreateMemoryContainerActionTests.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.mockito.Captor;
2626
import org.mockito.Mock;
2727
import org.mockito.MockitoAnnotations;
28+
import org.opensearch.OpenSearchStatusException;
2829
import org.opensearch.ResourceAlreadyExistsException;
2930
import org.opensearch.action.DocWriteResponse;
3031
import org.opensearch.action.admin.indices.create.CreateIndexRequest;
@@ -38,6 +39,7 @@
3839
import org.opensearch.commons.authuser.User;
3940
import org.opensearch.core.action.ActionListener;
4041
import org.opensearch.core.index.shard.ShardId;
42+
import org.opensearch.core.rest.RestStatus;
4143
import org.opensearch.ml.common.FunctionName;
4244
import org.opensearch.ml.common.MLIndex;
4345
import org.opensearch.ml.common.MLModel;
@@ -160,6 +162,7 @@ public void setup() {
160162

161163
// Setup ML feature settings
162164
when(mlFeatureEnabledSetting.isMultiTenancyEnabled()).thenReturn(false);
165+
when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); // Enable by default for tests
163166

164167
// Create action
165168
action = new TransportCreateMemoryContainerAction(
@@ -1386,6 +1389,23 @@ public void testDoExecuteWithSemanticStorageEnabledButNullEmbeddingModelId() thr
13861389
verify(mlModelManager, never()).getModel(anyString(), any());
13871390
}
13881391

1392+
public void testDoExecuteWithAgenticMemoryDisabled() throws InterruptedException {
1393+
// Disable agentic memory feature
1394+
when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(false);
1395+
1396+
// Execute
1397+
action.doExecute(task, request, actionListener);
1398+
1399+
// Verify failure response due to feature being disabled
1400+
verify(actionListener).onFailure(exceptionCaptor.capture());
1401+
Exception exception = exceptionCaptor.getValue();
1402+
assertNotNull(exception);
1403+
assertTrue(exception instanceof OpenSearchStatusException);
1404+
OpenSearchStatusException statusException = (OpenSearchStatusException) exception;
1405+
assertEquals(RestStatus.FORBIDDEN, statusException.status());
1406+
assertEquals("The Agentic Memory APIs are not enabled. To enable, please update the setting plugins.ml_commons.agentic_memory_enabled", exception.getMessage());
1407+
}
1408+
13891409
public void testDoExecuteWithSemanticStorageEnabledButNoLlmModel() throws InterruptedException {
13901410
// Create config with semantic storage enabled but no LLM model ID
13911411
MemoryStorageConfig configWithoutLlmModel = MemoryStorageConfig

0 commit comments

Comments
 (0)