Merge branch 'main' into intro-resource-permissions

Author: dhrubo-os

build.gradle

@@ -11,8 +11,8 @@ buildscript {
     ext {
         opensearch_group = "org.opensearch"
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")
-        opensearch_version = System.getProperty("opensearch.version", "3.0.0-SNAPSHOT")
-        buildVersionQualifier = System.getProperty("build.version_qualifier", "beta1")
+        opensearch_version = System.getProperty("opensearch.version", "3.1.0-SNAPSHOT")
+        buildVersionQualifier = System.getProperty("build.version_qualifier", "")
         asm_version = "9.7"
 
         // 2.0.0-rc1-SNAPSHOT -> 2.0.0.0-rc1-SNAPSHOT

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

@@ -8,7 +8,6 @@
 import java.util.Set;
 
 import org.opensearch.Version;
-import org.opensearch.common.settings.Setting;
 
 import com.google.common.collect.ImmutableSet;
 
@@ -46,6 +45,7 @@ public class CommonValue {
     public static final String ML_MEMORY_MESSAGE_INDEX = ".plugins-ml-memory-message";
     public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
     public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
+    public static final String MCP_SESSION_MANAGEMENT_INDEX = ".plugins-ml-mcp-session-management";
     public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
     public static final String TOOL_PARAMETERS_PREFIX = "tools.parameters.";
 
@@ -59,6 +59,7 @@ public class CommonValue {
     public static final String ML_AGENT_INDEX_MAPPING_PATH = "index-mappings/ml_agent.json";
     public static final String ML_MEMORY_META_INDEX_MAPPING_PATH = "index-mappings/ml_memory_meta.json";
     public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
+    public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
 
     // Calculate Versions independently of OpenSearch core version
     public static final Version VERSION_2_11_0 = Version.fromString("2.11.0");
@@ -93,16 +94,10 @@ public class CommonValue {
     public static final String MCP_TOOL_DESCRIPTION_FIELD = "description";
     public static final String MCP_TOOL_INPUT_SCHEMA_FIELD = "inputSchema";
     public static final String MCP_SYNC_CLIENT = "mcp_sync_client";
-    public static final String MCP_EXECUTOR_SERVICE = "mcp_executor_service";
     public static final String MCP_TOOLS_FIELD = "tools";
     public static final String MCP_CONNECTORS_FIELD = "mcp_connectors";
     public static final String MCP_CONNECTOR_ID_FIELD = "mcp_connector_id";
 
-    public static final Setting<Boolean> ML_COMMONS_MCP_FEATURE_ENABLED = Setting
-        .boolSetting("plugins.ml_commons.mcp_feature_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
-    public static final String ML_COMMONS_MCP_FEATURE_DISABLED_MESSAGE =
-        "The MCP feature is not enabled. To enable, please update the setting " + ML_COMMONS_MCP_FEATURE_ENABLED.getKey();
-
     // TOOL Constants
     public static final String TOOL_INPUT_SCHEMA_FIELD = "input_schema";
 }

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

@@ -13,6 +13,7 @@
 import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX_MAPPING_PATH;
+import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
@@ -38,7 +39,8 @@ public enum MLIndex {
     CONTROLLER(ML_CONTROLLER_INDEX, false, ML_CONTROLLER_INDEX_MAPPING_PATH),
     AGENT(ML_AGENT_INDEX, false, ML_AGENT_INDEX_MAPPING_PATH),
     MEMORY_META(ML_MEMORY_META_INDEX, false, ML_MEMORY_META_INDEX_MAPPING_PATH),
-    MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH);
+    MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH),
+    MCP_SESSION_MANAGEMENT(ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH);
 
     private final String indexName;
     // whether we use an alias for the index

common/src/main/java/org/opensearch/ml/common/ToolMetadata.java

@@ -7,7 +7,9 @@
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 
 import java.io.IOException;
+import java.util.Map;
 
+import org.opensearch.Version;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
@@ -25,6 +27,9 @@ public class ToolMetadata implements ToXContentObject, Writeable {
     public static final String TOOL_DESCRIPTION_FIELD = "description";
     public static final String TOOL_TYPE_FIELD = "type";
     public static final String TOOL_VERSION_FIELD = "version";
+    public static final String TOOL_ATTRIBUTES_FIELD = "attributes";
+
+    private static final Version MINIMUM_VERSION_FOR_TOOL_ATTRIBUTES = Version.V_3_0_0;
 
     @Getter
     private String name;
@@ -34,27 +39,43 @@ public class ToolMetadata implements ToXContentObject, Writeable {
     private String type;
     @Getter
     private String version;
+    @Getter
+    private Map<String, Object> attributes;
 
     @Builder(toBuilder = true)
-    public ToolMetadata(String name, String description, String type, String version) {
+    public ToolMetadata(String name, String description, String type, String version, Map<String, Object> attributes) {
         this.name = name;
         this.description = description;
         this.type = type;
         this.version = version;
+        this.attributes = attributes;
     }
 
     public ToolMetadata(StreamInput input) throws IOException {
+        Version byteStreamVersion = input.getVersion();
         name = input.readString();
         description = input.readString();
         type = input.readString();
         version = input.readOptionalString();
+        if (byteStreamVersion.onOrAfter(MINIMUM_VERSION_FOR_TOOL_ATTRIBUTES) && input.readBoolean()) {
+            attributes = input.readMap(StreamInput::readString, StreamInput::readGenericValue);
+        }
     }
 
     public void writeTo(StreamOutput output) throws IOException {
+        Version byteStreamVersion = output.getVersion();
         output.writeString(name);
         output.writeString(description);
         output.writeString(type);
         output.writeOptionalString(version);
+        if (byteStreamVersion.onOrAfter(MINIMUM_VERSION_FOR_TOOL_ATTRIBUTES)) {
+            if (attributes != null) {
+                output.writeBoolean(true);
+                output.writeMap(attributes, StreamOutput::writeString, StreamOutput::writeGenericValue);
+            } else {
+                output.writeBoolean(false);
+            }
+        }
     }
 
     @Override
@@ -70,6 +91,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
             builder.field(TOOL_TYPE_FIELD, type);
         }
         builder.field(TOOL_VERSION_FIELD, version != null ? version : "undefined");
+        if (attributes != null) {
+            builder.field(TOOL_ATTRIBUTES_FIELD, attributes);
+        }
         builder.endObject();
         return builder;
     }
@@ -79,6 +103,7 @@ public static ToolMetadata parse(XContentParser parser) throws IOException {
         String description = null;
         String type = null;
         String version = null;
+        Map<String, Object> attributes = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -97,12 +122,14 @@ public static ToolMetadata parse(XContentParser parser) throws IOException {
                     break;
                 case TOOL_VERSION_FIELD:
                     version = parser.text();
+                case TOOL_ATTRIBUTES_FIELD:
+                    attributes = parser.map();
                 default:
                     parser.skipChildren();
                     break;
             }
         }
-        return ToolMetadata.builder().name(name).description(description).type(type).version(version).build();
+        return ToolMetadata.builder().name(name).description(description).type(type).version(version).attributes(attributes).build();
     }
 
     public static ToolMetadata fromStream(StreamInput in) throws IOException {

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

@@ -16,7 +16,6 @@
 import org.opensearch.common.settings.Setting;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.core.common.unit.ByteSizeValue;
-import org.opensearch.ml.common.CommonValue;
 
 import com.google.common.collect.ImmutableList;
 
@@ -217,7 +216,15 @@ private MLCommonsSettings() {}
     public static final Setting<Boolean> ML_COMMONS_MEMORY_FEATURE_ENABLED = Setting
         .boolSetting("plugins.ml_commons.memory_feature_enabled", true, Setting.Property.NodeScope, Setting.Property.Dynamic);
 
-    public static final Setting<Boolean> ML_COMMONS_MCP_FEATURE_ENABLED = CommonValue.ML_COMMONS_MCP_FEATURE_ENABLED;
+    public static final Setting<Boolean> ML_COMMONS_MCP_CONNECTOR_ENABLED = Setting
+        .boolSetting("plugins.ml_commons.mcp_connector_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
+    public static final String ML_COMMONS_MCP_CONNECTOR_DISABLED_MESSAGE =
+        "The MCP connector is not enabled. To enable, please update the setting " + ML_COMMONS_MCP_CONNECTOR_ENABLED.getKey();
+
+    public static final Setting<Boolean> ML_COMMONS_MCP_SERVER_ENABLED = Setting
+        .boolSetting("plugins.ml_commons.mcp_server_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
+    public static final String ML_COMMONS_MCP_SERVER_DISABLED_MESSAGE =
+        "The MCP server is not enabled. To enable, please update the setting " + ML_COMMONS_MCP_SERVER_ENABLED.getKey();
 
     // Feature flag for enabling search processors for Retrieval Augmented Generation using OpenSearch and Remote Inference.
     public static final Setting<Boolean> ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED = Setting

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.mcpserver.action;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
+
+public class MLMcpMessageAction extends ActionType<AcknowledgedResponse> {
+    public static MLMcpMessageAction INSTANCE = new MLMcpMessageAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/mcp/message";
+
+    private MLMcpMessageAction() {
+        super(NAME, AcknowledgedResponse::new);
+    }
+
+}

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.mcpserver.action;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
+
+public class MLMcpMessageDispatchAction extends ActionType<AcknowledgedResponse> {
+    public static MLMcpMessageDispatchAction INSTANCE = new MLMcpMessageDispatchAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/mcp/message/dispatch";
+
+    private MLMcpMessageDispatchAction() {
+        super(NAME, AcknowledgedResponse::new);
+    }
+
+}

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.mcpserver.action;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.ml.common.transport.mcpserver.responses.register.MLMcpRegisterNodesResponse;
+
+public class MLMcpToolsRegisterOnNodesAction extends ActionType<MLMcpRegisterNodesResponse> {
+    public static MLMcpToolsRegisterOnNodesAction INSTANCE = new MLMcpToolsRegisterOnNodesAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/register_on_nodes";
+
+    private MLMcpToolsRegisterOnNodesAction() {
+        super(NAME, MLMcpRegisterNodesResponse::new);
+    }
+
+}

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.mcpserver.action;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.ml.common.transport.mcpserver.responses.remove.MLMcpRemoveNodesResponse;
+
+public class MLMcpToolsRemoveOnNodesAction extends ActionType<MLMcpRemoveNodesResponse> {
+    public static MLMcpToolsRemoveOnNodesAction INSTANCE = new MLMcpToolsRemoveOnNodesAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/remove_on_nodes";
+
+    private MLMcpToolsRemoveOnNodesAction() {
+        super(NAME, MLMcpRemoveNodesResponse::new);
+    }
+
+}

common/src/main/java/org/opensearch/ml/common/transport/mcpserver/requests/message/MLMcpMessageRequest.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.mcpserver.requests.message;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import org.apache.commons.lang3.StringUtils;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.core.common.io.stream.InputStreamStreamInput;
+import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.transport.TransportRequest;
+
+import lombok.Builder;
+import lombok.Getter;
+
+@Getter
+public class MLMcpMessageRequest extends ActionRequest {
+
+    private final String nodeId;
+
+    private final String sessionId;
+
+    private final String requestBody;
+
+    public MLMcpMessageRequest(StreamInput in) throws IOException {
+        super(in);
+        this.nodeId = in.readString();
+        this.sessionId = in.readString();
+        this.requestBody = in.readString();
+        if (StringUtils.isEmpty(nodeId) || StringUtils.isEmpty(sessionId) || StringUtils.isEmpty(requestBody)) {
+            throw new IllegalStateException("nodeId, sessionId and requestBody must not be null");
+        }
+    }
+
+    @Builder
+    public MLMcpMessageRequest(String nodeId, String sessionId, String requestBody) {
+        super();
+        if (StringUtils.isEmpty(nodeId) || StringUtils.isEmpty(sessionId) || StringUtils.isEmpty(requestBody)) {
+            throw new IllegalStateException("nodeId, sessionId and requestBody must not be null");
+        }
+        this.nodeId = nodeId;
+        this.sessionId = sessionId;
+        this.requestBody = requestBody;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeString(nodeId);
+        out.writeString(sessionId);
+        out.writeString(requestBody);
+    }
+
+    public static MLMcpMessageRequest fromActionRequest(TransportRequest actionRequest) {
+        if (actionRequest instanceof MLMcpMessageRequest) {
+            return (MLMcpMessageRequest) actionRequest;
+        }
+
+        try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
+            actionRequest.writeTo(osso);
+            try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
+                return new MLMcpMessageRequest(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("Failed to parse ActionRequest into MLMcpMessageRequest", e);
+        }
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+}