fix connector executor: restore context

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

@@ -54,6 +54,7 @@ public class CommonValue {
     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 String MCP_TOOLS_INDEX = ".plugins-ml-mcp-tools";
+    public static final String ML_CONTEXT_MANAGEMENT_TEMPLATES_INDEX = ".plugins-ml-context-management-templates";
     // index created in 3.1 to track all ml jobs created via job scheduler
     public static final String ML_JOBS_INDEX = ".plugins-ml-jobs";
     public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
@@ -76,6 +77,7 @@ public class CommonValue {
     public static final String ML_LONG_MEMORY_HISTORY_INDEX_MAPPING_PATH = "index-mappings/ml_memory_long_term_history.json";
     public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
     public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
+    public static final String ML_CONTEXT_MANAGEMENT_TEMPLATES_INDEX_MAPPING_PATH = "index-mappings/ml_context_management_templates.json";
     public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";
     public static final String ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_config.json";
     public static final String ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_storage.json";
@@ -97,6 +99,7 @@ public class CommonValue {
     public static final Version VERSION_3_1_0 = Version.fromString("3.1.0");
     public static final Version VERSION_3_2_0 = Version.fromString("3.2.0");
     public static final Version VERSION_3_3_0 = Version.fromString("3.3.0");
+    public static final Version VERSION_3_4_0 = Version.fromString("3.4.0");
 
     // Connector Constants
     public static final String NAME_FIELD = "name";
@@ -113,6 +116,7 @@ public class CommonValue {
     public static final String CLIENT_CONFIG_FIELD = "client_config";
     public static final String URL_FIELD = "url";
     public static final String HEADERS_FIELD = "headers";
+    public static final String CONNECTOR_ACTION_FIELD = "connector_action";
 
     // MCP Constants
     public static final String MCP_TOOL_NAME_FIELD = "name";

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

@@ -0,0 +1,25 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common;
+
+import java.util.Locale;
+
+public enum MLMemoryType {
+    CONVERSATION_INDEX,
+    AGENTIC_MEMORY,
+    REMOTE_AGENTIC_MEMORY;
+
+    public static MLMemoryType from(String value) {
+        if (value != null) {
+            try {
+                return MLMemoryType.valueOf(value.toUpperCase(Locale.ROOT));
+            } catch (Exception e) {
+                throw new IllegalArgumentException("Wrong Memory type");
+            }
+        }
+        return null;
+    }
+}

common/src/main/java/org/opensearch/ml/common/agent/MLAgent.java

@@ -30,6 +30,7 @@
 import org.opensearch.ml.common.CommonValue;
 import org.opensearch.ml.common.MLAgentType;
 import org.opensearch.ml.common.MLModel;
+import org.opensearch.ml.common.contextmanager.ContextManagementTemplate;
 import org.opensearch.telemetry.metrics.tags.Tags;
 
 import lombok.Builder;
@@ -51,13 +52,16 @@ public class MLAgent implements ToXContentObject, Writeable {
     public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
     public static final String APP_TYPE_FIELD = "app_type";
     public static final String IS_HIDDEN_FIELD = "is_hidden";
+    public static final String CONTEXT_MANAGEMENT_NAME_FIELD = "context_management_name";
+    public static final String CONTEXT_MANAGEMENT_FIELD = "context_management";
     private static final String LLM_INTERFACE_FIELD = "_llm_interface";
     private static final String TAG_VALUE_UNKNOWN = "unknown";
     private static final String TAG_MEMORY_TYPE = "memory_type";
 
     public static final int AGENT_NAME_MAX_LENGTH = 128;
 
     private static final Version MINIMAL_SUPPORTED_VERSION_FOR_HIDDEN_AGENT = CommonValue.VERSION_2_13_0;
+    private static final Version MINIMAL_SUPPORTED_VERSION_FOR_CONTEXT_MANAGEMENT = CommonValue.VERSION_3_3_0;
 
     private String name;
     private String type;
@@ -71,6 +75,8 @@ public class MLAgent implements ToXContentObject, Writeable {
     private Instant lastUpdateTime;
     private String appType;
     private Boolean isHidden;
+    private String contextManagementName;
+    private ContextManagementTemplate contextManagement;
     private final String tenantId;
 
     @Builder(toBuilder = true)
@@ -86,6 +92,8 @@ public MLAgent(
         Instant lastUpdateTime,
         String appType,
         Boolean isHidden,
+        String contextManagementName,
+        ContextManagementTemplate contextManagement,
         String tenantId
     ) {
         this.name = name;
@@ -100,6 +108,8 @@ public MLAgent(
         this.appType = appType;
         // is_hidden field isn't going to be set by user. It will be set by the code.
         this.isHidden = isHidden;
+        this.contextManagementName = contextManagementName;
+        this.contextManagement = contextManagement;
         this.tenantId = tenantId;
         validate();
     }
@@ -128,6 +138,17 @@ private void validate() {
                 }
             }
         }
+        validateContextManagement();
+    }
+
+    private void validateContextManagement() {
+        if (contextManagementName != null && contextManagement != null) {
+            throw new IllegalArgumentException("Cannot specify both context_management_name and context_management");
+        }
+
+        if (contextManagement != null && !contextManagement.isValid()) {
+            throw new IllegalArgumentException("Invalid context management configuration");
+        }
     }
 
     private void validateMLAgentType(String agentType) {
@@ -171,6 +192,12 @@ public MLAgent(StreamInput input) throws IOException {
         if (streamInputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_HIDDEN_AGENT)) {
             isHidden = input.readOptionalBoolean();
         }
+        if (streamInputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_CONTEXT_MANAGEMENT)) {
+            contextManagementName = input.readOptionalString();
+            if (input.readBoolean()) {
+                contextManagement = new ContextManagementTemplate(input);
+            }
+        }
         this.tenantId = streamInputVersion.onOrAfter(VERSION_2_19_0) ? input.readOptionalString() : null;
         validate();
     }
@@ -214,6 +241,15 @@ public void writeTo(StreamOutput out) throws IOException {
         if (streamOutputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_HIDDEN_AGENT)) {
             out.writeOptionalBoolean(isHidden);
         }
+        if (streamOutputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_CONTEXT_MANAGEMENT)) {
+            out.writeOptionalString(contextManagementName);
+            if (contextManagement != null) {
+                out.writeBoolean(true);
+                contextManagement.writeTo(out);
+            } else {
+                out.writeBoolean(false);
+            }
+        }
         if (streamOutputVersion.onOrAfter(VERSION_2_19_0)) {
             out.writeOptionalString(tenantId);
         }
@@ -256,6 +292,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (isHidden != null) {
             builder.field(MLModel.IS_HIDDEN_FIELD, isHidden);
         }
+        if (contextManagementName != null) {
+            builder.field(CONTEXT_MANAGEMENT_NAME_FIELD, contextManagementName);
+        }
+        if (contextManagement != null) {
+            builder.field(CONTEXT_MANAGEMENT_FIELD, contextManagement);
+        }
         if (tenantId != null) {
             builder.field(TENANT_ID_FIELD, tenantId);
         }
@@ -283,6 +325,8 @@ private static MLAgent parseCommonFields(XContentParser parser, boolean parseHid
         Instant lastUpdateTime = null;
         String appType = null;
         boolean isHidden = false;
+        String contextManagementName = null;
+        ContextManagementTemplate contextManagement = null;
         String tenantId = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
@@ -329,6 +373,12 @@ private static MLAgent parseCommonFields(XContentParser parser, boolean parseHid
                     if (parseHidden)
                         isHidden = parser.booleanValue();
                     break;
+                case CONTEXT_MANAGEMENT_NAME_FIELD:
+                    contextManagementName = parser.text();
+                    break;
+                case CONTEXT_MANAGEMENT_FIELD:
+                    contextManagement = ContextManagementTemplate.parse(parser);
+                    break;
                 case TENANT_ID_FIELD:
                     tenantId = parser.textOrNull();
                     break;
@@ -351,6 +401,8 @@ private static MLAgent parseCommonFields(XContentParser parser, boolean parseHid
             .lastUpdateTime(lastUpdateTime)
             .appType(appType)
             .isHidden(isHidden)
+            .contextManagementName(contextManagementName)
+            .contextManagement(contextManagement)
             .tenantId(tenantId)
             .build();
     }
@@ -384,4 +436,47 @@ public Tags getTags() {
 
         return tags;
     }
+
+    /**
+     * Check if this agent has context management configuration
+     * @return true if agent has either context management name or inline configuration
+     */
+    public boolean hasContextManagement() {
+        return contextManagementName != null || contextManagement != null;
+    }
+
+    /**
+     * Get the effective context management configuration for this agent.
+     * This method prioritizes inline configuration over template reference.
+     * Note: Template resolution requires external service call and should be handled by the caller.
+     * 
+     * @return the inline context management configuration, or null if using template reference or no configuration
+     */
+    public ContextManagementTemplate getInlineContextManagement() {
+        return contextManagement;
+    }
+
+    /**
+     * Check if this agent uses a context management template reference
+     * @return true if agent references a context management template by name
+     */
+    public boolean hasContextManagementTemplate() {
+        return contextManagementName != null;
+    }
+
+    /**
+     * Check if this agent has inline context management configuration
+     * @return true if agent has inline context management configuration
+     */
+    public boolean hasInlineContextManagement() {
+        return contextManagement != null;
+    }
+
+    /**
+     * Get the context management template name if this agent references one
+     * @return the template name, or null if no template reference
+     */
+    public String getContextManagementTemplateName() {
+        return contextManagementName;
+    }
 }

common/src/main/java/org/opensearch/ml/common/agent/MLMemorySpec.java

@@ -26,32 +26,37 @@ public class MLMemorySpec implements ToXContentObject {
     public static final String MEMORY_TYPE_FIELD = "type";
     public static final String WINDOW_SIZE_FIELD = "window_size";
     public static final String SESSION_ID_FIELD = "session_id";
+    public static final String MEMORY_CONTAINER_ID_FIELD = "memory_container_id";
 
     private String type;
     @Setter
     private String sessionId;
     private Integer windowSize;
+    private String memoryContainerId;
 
     @Builder(toBuilder = true)
-    public MLMemorySpec(String type, String sessionId, Integer windowSize) {
+    public MLMemorySpec(String type, String sessionId, Integer windowSize, String memoryContainerId) {
         if (type == null) {
             throw new IllegalArgumentException("agent name is null");
         }
         this.type = type;
         this.sessionId = sessionId;
         this.windowSize = windowSize;
+        this.memoryContainerId = memoryContainerId;
     }
 
     public MLMemorySpec(StreamInput input) throws IOException {
         type = input.readString();
         sessionId = input.readOptionalString();
         windowSize = input.readOptionalInt();
+        memoryContainerId = input.readOptionalString();
     }
 
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(type);
         out.writeOptionalString(sessionId);
         out.writeOptionalInt(windowSize);
+        out.writeOptionalString(memoryContainerId);
     }
 
     @Override
@@ -64,6 +69,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (sessionId != null) {
             builder.field(SESSION_ID_FIELD, sessionId);
         }
+        if (memoryContainerId != null) {
+            builder.field(MEMORY_CONTAINER_ID_FIELD, memoryContainerId);
+        }
         builder.endObject();
         return builder;
     }
@@ -72,6 +80,7 @@ public static MLMemorySpec parse(XContentParser parser) throws IOException {
         String type = null;
         String sessionId = null;
         Integer windowSize = null;
+        String memoryContainerId = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -88,12 +97,15 @@ public static MLMemorySpec parse(XContentParser parser) throws IOException {
                 case WINDOW_SIZE_FIELD:
                     windowSize = parser.intValue();
                     break;
+                case MEMORY_CONTAINER_ID_FIELD:
+                    memoryContainerId = parser.text();
+                    break;
                 default:
                     parser.skipChildren();
                     break;
             }
         }
-        return MLMemorySpec.builder().type(type).sessionId(sessionId).windowSize(windowSize).build();
+        return MLMemorySpec.builder().type(type).sessionId(sessionId).windowSize(windowSize).memoryContainerId(memoryContainerId).build();
     }
 
     public static MLMemorySpec fromStream(StreamInput in) throws IOException {

common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java

@@ -121,8 +121,11 @@ public <T> void parseResponse(T response, List<ModelTensor> modelTensors, boolea
 
     @Override
     public Optional<ConnectorAction> findAction(String action) {
-        if (actions != null) {
-            return actions.stream().filter(a -> a.getActionType().name().equalsIgnoreCase(action)).findFirst();
+        if (actions != null && action != null) {
+            if (ConnectorAction.ActionType.isValidAction(action)) {
+                return actions.stream().filter(a -> a.getActionType().name().equalsIgnoreCase(action)).findFirst();
+            }
+            return actions.stream().filter(a -> action.equals(a.getName())).findFirst();
         }
         return Optional.empty();
     }