Skip to content

Commit 2df984d

Browse files
zane-neozane-neo
authored
Support persisting MCP tools in system index (#3874)
* Support MCP tools configuration persistence Signed-off-by: zane-neo <zaniu@amazon.com> * Support persisting mcp tools configuration in system index Signed-off-by: zane-neo <zaniu@amazon.com> * format code Signed-off-by: zane-neo <zaniu@amazon.com> * Address comments Signed-off-by: zane-neo <zaniu@amazon.com> * add ignore code coverage Signed-off-by: zane-neo <zaniu@amazon.com> * fix failure UTs Signed-off-by: zane-neo <zaniu@amazon.com> * exclude files that unable to cover Signed-off-by: zane-neo <zaniu@amazon.com> * addressing comments Signed-off-by: zane-neo <zaniu@amazon.com> * Addressing comments Signed-off-by: zane-neo <zaniu@amazon.com> * address comments Signed-off-by: zane-neo <zaniu@amazon.com> * fix failure UT Signed-off-by: zane-neo <zaniu@amazon.com> * fix failure UT Signed-off-by: zane-neo <zaniu@amazon.com> --------- Signed-off-by: zane-neo <zaniu@amazon.com>
1 parent a866cec commit 2df984d

File tree

83 files changed

+6009
-844
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

83 files changed

+6009
-844
lines changed

‎common/src/main/java/org/opensearch/ml/common/CommonValue.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public class CommonValue {
4646
public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
4747
public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
4848
public static final String MCP_SESSION_MANAGEMENT_INDEX = ".plugins-ml-mcp-session-management";
49+
public static final String MCP_TOOLS_INDEX = ".plugins-ml-mcp-tools";
4950
public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
5051
public static final String TOOL_PARAMETERS_PREFIX = "tools.parameters.";
5152

@@ -60,6 +61,7 @@ public class CommonValue {
6061
public static final String ML_MEMORY_META_INDEX_MAPPING_PATH = "index-mappings/ml_memory_meta.json";
6162
public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
6263
public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
64+
public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
6365

6466
// Calculate Versions independently of OpenSearch core version
6567
public static final Version VERSION_2_11_0 = Version.fromString("2.11.0");

‎common/src/main/java/org/opensearch/ml/common/MLIndex.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
package org.opensearch.ml.common;
77

8+
import static org.opensearch.ml.common.CommonValue.MCP_SESSION_MANAGEMENT_INDEX;
9+
import static org.opensearch.ml.common.CommonValue.MCP_TOOLS_INDEX;
810
import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
911
import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX_MAPPING_PATH;
1012
import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;
@@ -14,6 +16,7 @@
1416
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
1517
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX_MAPPING_PATH;
1618
import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
19+
import static org.opensearch.ml.common.CommonValue.ML_MCP_TOOLS_INDEX_MAPPING_PATH;
1720
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
1821
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH;
1922
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
@@ -40,7 +43,8 @@ public enum MLIndex {
4043
AGENT(ML_AGENT_INDEX, false, ML_AGENT_INDEX_MAPPING_PATH),
4144
MEMORY_META(ML_MEMORY_META_INDEX, false, ML_MEMORY_META_INDEX_MAPPING_PATH),
4245
MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH),
43-
MCP_SESSION_MANAGEMENT(ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH);
46+
MCP_SESSION_MANAGEMENT(MCP_SESSION_MANAGEMENT_INDEX, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH),
47+
MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH);
4448

4549
private final String indexName;
4650
// whether we use an alias for the index

‎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
@@ -11,6 +11,7 @@
1111
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED;
1212
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED;
1313
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_LOCAL_MODEL_ENABLED;
14+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MCP_SERVER_ENABLED;
1415
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MULTI_TENANCY_ENABLED;
1516
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED;
1617
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED;
@@ -40,6 +41,8 @@ public class MLFeatureEnabledSetting {
4041
// This is to identify if this node is in multi-tenancy or not.
4142
private volatile Boolean isMultiTenancyEnabled;
4243

44+
private volatile Boolean isMcpServerEnabled;
45+
4346
private final List<SettingsChangeListener> listeners = new ArrayList<>();
4447

4548
public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings) {
@@ -51,6 +54,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
5154
isBatchIngestionEnabled = ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED.get(settings);
5255
isBatchInferenceEnabled = ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED.get(settings);
5356
isMultiTenancyEnabled = ML_COMMONS_MULTI_TENANCY_ENABLED.get(settings);
57+
isMcpServerEnabled = ML_COMMONS_MCP_SERVER_ENABLED.get(settings);
5458

5559
clusterService
5660
.getClusterSettings()
@@ -69,6 +73,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
6973
clusterService
7074
.getClusterSettings()
7175
.addSettingsUpdateConsumer(ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED, it -> isBatchInferenceEnabled = it);
76+
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MCP_SERVER_ENABLED, it -> isMcpServerEnabled = it);
7277
}
7378

7479
/**
@@ -131,6 +136,14 @@ public boolean isMultiTenancyEnabled() {
131136
return isMultiTenancyEnabled;
132137
}
133138

139+
/**
140+
* Whether the mcp server feature is enabled. If disabled, MCP server APIs in ml-commons will be blocked.
141+
* @return where the MCP server feature is enabled.
142+
*/
143+
public boolean isMcpServerEnabled() {
144+
return isMcpServerEnabled;
145+
}
146+
134147
public void addListener(SettingsChangeListener listener) {
135148
listeners.add(listener);
136149
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsListAction.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.list.MLMcpToolsListResponse;
10+
11+
public class MLMcpToolsListAction extends ActionType<MLMcpToolsListResponse> {
12+
public static MLMcpToolsListAction INSTANCE = new MLMcpToolsListAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/list";
14+
15+
private MLMcpToolsListAction() {
16+
super(NAME, MLMcpToolsListResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsRegisterAction.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.register.MLMcpToolsRegisterNodesResponse;
10+
11+
public class MLMcpToolsRegisterAction extends ActionType<MLMcpToolsRegisterNodesResponse> {
12+
public static MLMcpToolsRegisterAction INSTANCE = new MLMcpToolsRegisterAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/register";
14+
15+
private MLMcpToolsRegisterAction() {
16+
super(NAME, MLMcpToolsRegisterNodesResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsRegisterOnNodesAction.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,14 @@
66
package org.opensearch.ml.common.transport.mcpserver.action;
77

88
import org.opensearch.action.ActionType;
9-
import org.opensearch.ml.common.transport.mcpserver.responses.register.MLMcpRegisterNodesResponse;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.register.MLMcpToolsRegisterNodesResponse;
1010

11-
public class MLMcpToolsRegisterOnNodesAction extends ActionType<MLMcpRegisterNodesResponse> {
11+
public class MLMcpToolsRegisterOnNodesAction extends ActionType<MLMcpToolsRegisterNodesResponse> {
1212
public static MLMcpToolsRegisterOnNodesAction INSTANCE = new MLMcpToolsRegisterOnNodesAction();
1313
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/register_on_nodes";
1414

1515
private MLMcpToolsRegisterOnNodesAction() {
16-
super(NAME, MLMcpRegisterNodesResponse::new);
16+
super(NAME, MLMcpToolsRegisterNodesResponse::new);
1717
}
1818

1919
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsRemoveAction.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.remove.MLMcpToolsRemoveNodesResponse;
10+
11+
public class MLMcpToolsRemoveAction extends ActionType<MLMcpToolsRemoveNodesResponse> {
12+
public static MLMcpToolsRemoveAction INSTANCE = new MLMcpToolsRemoveAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/remove";
14+
15+
private MLMcpToolsRemoveAction() {
16+
super(NAME, MLMcpToolsRemoveNodesResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsRemoveOnNodesAction.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,14 @@
66
package org.opensearch.ml.common.transport.mcpserver.action;
77

88
import org.opensearch.action.ActionType;
9-
import org.opensearch.ml.common.transport.mcpserver.responses.remove.MLMcpRemoveNodesResponse;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.remove.MLMcpToolsRemoveNodesResponse;
1010

11-
public class MLMcpToolsRemoveOnNodesAction extends ActionType<MLMcpRemoveNodesResponse> {
11+
public class MLMcpToolsRemoveOnNodesAction extends ActionType<MLMcpToolsRemoveNodesResponse> {
1212
public static MLMcpToolsRemoveOnNodesAction INSTANCE = new MLMcpToolsRemoveOnNodesAction();
1313
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/remove_on_nodes";
1414

1515
private MLMcpToolsRemoveOnNodesAction() {
16-
super(NAME, MLMcpRemoveNodesResponse::new);
16+
super(NAME, MLMcpToolsRemoveNodesResponse::new);
1717
}
1818

1919
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsUpdateAction.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.update.MLMcpToolsUpdateNodesResponse;
10+
11+
public class MLMcpToolsUpdateAction extends ActionType<MLMcpToolsUpdateNodesResponse> {
12+
public static MLMcpToolsUpdateAction INSTANCE = new MLMcpToolsUpdateAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/update";
14+
15+
private MLMcpToolsUpdateAction() {
16+
super(NAME, MLMcpToolsUpdateNodesResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsUpdateOnNodesAction.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.update.MLMcpToolsUpdateNodesResponse;
10+
11+
public class MLMcpToolsUpdateOnNodesAction extends ActionType<MLMcpToolsUpdateNodesResponse> {
12+
public static MLMcpToolsUpdateOnNodesAction INSTANCE = new MLMcpToolsUpdateOnNodesAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/update_on_nodes";
14+
15+
private MLMcpToolsUpdateOnNodesAction() {
16+
super(NAME, MLMcpToolsUpdateNodesResponse::new);
17+
}
18+
19+
}

0 commit comments

Comments
 (0)