opensearch-project
diff --git a/‎common/src/main/java/org/opensearch/ml/common/CommonValue.java‎
Lines changed: 2 additions & 0 deletions b/‎common/src/main/java/org/opensearch/ml/common/CommonValue.java‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎common/src/main/java/org/opensearch/ml/common/agent/MLAgent.java‎
Lines changed: 95 additions & 0 deletions b/‎common/src/main/java/org/opensearch/ml/common/agent/MLAgent.java‎
Lines changed: 95 additions & 0 deletions
diff --git a/‎common/src/main/java/org/opensearch/ml/common/contextmanager/ActivationRule.java‎
Lines changed: 26 additions & 0 deletions b/‎common/src/main/java/org/opensearch/ml/common/contextmanager/ActivationRule.java‎
Lines changed: 26 additions & 0 deletions
diff --git a/‎common/src/main/java/org/opensearch/ml/common/contextmanager/ActivationRuleFactory.java‎
Lines changed: 148 additions & 0 deletions b/‎common/src/main/java/org/opensearch/ml/common/contextmanager/ActivationRuleFactory.java‎
Lines changed: 148 additions & 0 deletions
@@ -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";
 
@@ -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;
+    }
 }
@@ -0,0 +1,26 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.contextmanager;
+
+/**
+ * Interface for activation rules that determine when a context manager should execute.
+ * Activation rules evaluate runtime conditions based on the current context state.
+ */
+public interface ActivationRule {
+
+    /**
+     * Evaluate whether the activation condition is met.
+     * @param context the current context state
+     * @return true if the condition is met and the manager should activate, false otherwise
+     */
+    boolean evaluate(ContextManagerContext context);
+
+    /**
+     * Get a description of this activation rule for logging and debugging.
+     * @return a human-readable description of the rule
+     */
+    String getDescription();
+}
@@ -0,0 +1,148 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.contextmanager;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import lombok.extern.log4j.Log4j2;
+
+/**
+ * Factory class for creating activation rules from configuration.
+ * Supports creating rules from configuration maps and combining multiple rules.
+ */
+@Log4j2
+public class ActivationRuleFactory {
+
+    public static final String TOKENS_EXCEED_KEY = "tokens_exceed";
+    public static final String MESSAGE_COUNT_EXCEED_KEY = "message_count_exceed";
+
+    /**
+     * Create activation rules from a configuration map.
+     * @param activationConfig the configuration map containing rule definitions
+     * @return a list of activation rules, or empty list if no valid rules found
+     */
+    public static List<ActivationRule> createRules(Map<String, Object> activationConfig) {
+        List<ActivationRule> rules = new ArrayList<>();
+
+        if (activationConfig == null || activationConfig.isEmpty()) {
+            return rules;
+        }
+
+        // Create tokens_exceed rule
+        if (activationConfig.containsKey(TOKENS_EXCEED_KEY)) {
+            try {
+                Object tokenValue = activationConfig.get(TOKENS_EXCEED_KEY);
+                int tokenThreshold = parseIntegerValue(tokenValue, TOKENS_EXCEED_KEY);
+                if (tokenThreshold > 0) {
+                    rules.add(new TokensExceedRule(tokenThreshold));
+                    log.debug("Created TokensExceedRule with threshold: {}", tokenThreshold);
+                } else {
+                    throw new IllegalArgumentException("Invalid token threshold value: " + tokenValue + ". Must be positive integer.");
+                }
+            } catch (Exception e) {
+                log.error("Failed to create TokensExceedRule: {}", e.getMessage());
+            }
+        }
+
+        // Create message_count_exceed rule
+        if (activationConfig.containsKey(MESSAGE_COUNT_EXCEED_KEY)) {
+            try {
+                Object messageValue = activationConfig.get(MESSAGE_COUNT_EXCEED_KEY);
+                int messageThreshold = parseIntegerValue(messageValue, MESSAGE_COUNT_EXCEED_KEY);
+                if (messageThreshold > 0) {
+                    rules.add(new MessageCountExceedRule(messageThreshold));
+                    log.debug("Created MessageCountExceedRule with threshold: {}", messageThreshold);
+                } else {
+                    throw new IllegalArgumentException(
+                        "Invalid message count threshold value: " + messageValue + ". Must be positive integer."
+                    );
+                }
+            } catch (Exception e) {
+                log.error("Failed to create MessageCountExceedRule: {}", e.getMessage());
+            }
+        }
+
+        return rules;
+    }
+
+    /**
+     * Create a composite rule that requires ALL rules to be satisfied (AND logic).
+     * @param rules the list of rules to combine
+     * @return a composite rule, or null if the list is empty
+     */
+    public static ActivationRule createCompositeRule(List<ActivationRule> rules) {
+        if (rules == null || rules.isEmpty()) {
+            return null;
+        }
+
+        if (rules.size() == 1) {
+            return rules.get(0);
+        }
+
+        return new CompositeActivationRule(rules);
+    }
+
+    /**
+     * Parse an integer value from configuration, handling various input types.
+     * @param value the value to parse
+     * @param fieldName the field name for error reporting
+     * @return the parsed integer value
+     * @throws IllegalArgumentException if the value cannot be parsed
+     */
+    private static int parseIntegerValue(Object value, String fieldName) {
+        if (value instanceof Integer) {
+            return (Integer) value;
+        } else if (value instanceof Number) {
+            return ((Number) value).intValue();
+        } else if (value instanceof String) {
+            try {
+                return Integer.parseInt((String) value);
+            } catch (NumberFormatException e) {
+                throw new IllegalArgumentException("Invalid integer value for " + fieldName + ": " + value);
+            }
+        } else {
+            throw new IllegalArgumentException("Unsupported value type for " + fieldName + ": " + value.getClass().getSimpleName());
+        }
+    }
+
+    /**
+     * Composite activation rule that implements AND logic for multiple rules.
+     */
+    private static class CompositeActivationRule implements ActivationRule {
+        private final List<ActivationRule> rules;
+
+        public CompositeActivationRule(List<ActivationRule> rules) {
+            this.rules = new ArrayList<>(rules);
+        }
+
+        @Override
+        public boolean evaluate(ContextManagerContext context) {
+            // All rules must evaluate to true (AND logic)
+            for (ActivationRule rule : rules) {
+                if (!rule.evaluate(context)) {
+                    return false;
+                }
+            }
+            return true;
+        }
+
+        @Override
+        public String getDescription() {
+            StringBuilder sb = new StringBuilder();
+            sb.append("composite_rule: [");
+            for (int i = 0; i < rules.size(); i++) {
+                if (i > 0) {
+                    sb.append(" AND ");
+                }
+                sb.append(rules.get(i).getDescription());
+            }
+            sb.append("]");
+            return sb.toString();
+        }
+    }
+}