Allow context management inline create in register agent without storing in index (#4403)

* allow inline create context management without storing in agent register

Signed-off-by: Mingshi Liu <mingshl@amazon.com>

* make ML_COMMONS_MULTI_TENANCY_ENABLED default is false

Signed-off-by: Mingshi Liu <mingshl@amazon.com>

---------

Signed-off-by: Mingshi Liu <mingshl@amazon.com>

Author: mingshl

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

@@ -464,6 +464,14 @@ 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

common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java

@@ -47,10 +47,16 @@ public ContextManagerHookProvider(List<ContextManager> contextManagers) {
      */
     @Override
     public void registerHooks(HookRegistry registry) {
-        // Register callbacks for each hook type
-        registry.addCallback(PreLLMEvent.class, this::handlePreLLM);
-        registry.addCallback(EnhancedPostToolEvent.class, this::handlePostTool);
-        registry.addCallback(PostMemoryEvent.class, this::handlePostMemory);
+        // Only register callbacks for hooks that have managers configured
+        if (hookToManagersMap.containsKey("PRE_LLM")) {
+            registry.addCallback(PreLLMEvent.class, this::handlePreLLM);
+        }
+        if (hookToManagersMap.containsKey("POST_TOOL")) {
+            registry.addCallback(EnhancedPostToolEvent.class, this::handlePostTool);
+        }
+        if (hookToManagersMap.containsKey("POST_MEMORY")) {
+            registry.addCallback(PostMemoryEvent.class, this::handlePostMemory);
+        }
 
         log.info("Registered context manager hooks for {} managers", contextManagers.size());
     }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java

@@ -484,6 +484,20 @@ private void processContextManagement(MLAgent mlAgent, HookRegistry hookRegistry
                 return; // Runtime parameter takes precedence, let MLExecuteTaskRunner handle it
             }
 
+            // Check if already processed to avoid duplicate registrations
+            if ("true".equals(inputDataSet.getParameters().get("context_management_processed"))) {
+                log.debug("Context management already processed for this execution, skipping");
+                return;
+            }
+
+            // Check if HookRegistry already has callbacks (from previous runtime setup)
+            // Don't override with inline configuration if runtime config is already active
+            if (hookRegistry.getCallbackCount(org.opensearch.ml.common.hooks.EnhancedPostToolEvent.class) > 0
+                || hookRegistry.getCallbackCount(org.opensearch.ml.common.hooks.PreLLMEvent.class) > 0) {
+                log.info("HookRegistry already has active configuration, skipping inline context management");
+                return;
+            }
+
             ContextManagementTemplate template = null;
             String templateName = null;
 

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java

@@ -484,6 +484,7 @@ private void executePlanningLoop(
                     .build();
 
                 // Pass hookRegistry to internal agent execution
+                // TODO need to check if the agentInput already have the hookResgistry?
                 agentInput.setHookRegistry(hookRegistry);
 
                 MLExecuteTaskRequest executeRequest = new MLExecuteTaskRequest(FunctionName.AGENT, agentInput);

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java

@@ -301,4 +301,21 @@ private MLAgent createTestAgent(String type) {
             .appType("test-app")
             .build();
     }
+
+    @Test
+    public void testContextManagementProcessedFlagPreventsReprocessing() {
+        // Test that the context_management_processed flag prevents duplicate processing
+        Map<String, String> parameters = new HashMap<>();
+
+        // First check - should allow processing
+        boolean shouldProcess1 = !"true".equals(parameters.get("context_management_processed"));
+        assertTrue("First call should allow processing", shouldProcess1);
+
+        // Mark as processed (simulating what the method does)
+        parameters.put("context_management_processed", "true");
+
+        // Second check - should prevent processing
+        boolean shouldProcess2 = !"true".equals(parameters.get("context_management_processed"));
+        assertFalse("Second call should prevent processing", shouldProcess2);
+    }
 }

plugin/src/test/java/org/opensearch/ml/action/contextmanagement/ContextManagerFactoryTests.java

@@ -140,4 +140,28 @@ public void testCreateContextManager_EmptyType() {
             assertTrue(e.getMessage().contains("Unsupported context manager type"));
         }
     }
+
+    @Test
+    public void testContextManagerHookProvider_SelectiveRegistration() {
+        // Test that ContextManagerHookProvider only registers hooks for configured managers
+        java.util.Map<String, java.util.List<String>> hookToManagersMap = new java.util.HashMap<>();
+
+        // Test 1: Only POST_TOOL configured
+        hookToManagersMap.put("POST_TOOL", java.util.Arrays.asList("ToolsOutputTruncateManager"));
+
+        // Simulate the registration logic
+        java.util.Set<String> registeredHooks = new java.util.HashSet<>();
+        if (hookToManagersMap.containsKey("PRE_LLM")) {
+            registeredHooks.add("PRE_LLM");
+        }
+        if (hookToManagersMap.containsKey("POST_TOOL")) {
+            registeredHooks.add("POST_TOOL");
+        }
+        if (hookToManagersMap.containsKey("POST_MEMORY")) {
+            registeredHooks.add("POST_MEMORY");
+        }
+
+        // Assert only POST_TOOL is registered
+        assertTrue("Should only register POST_TOOL hook", registeredHooks.size() == 1 && registeredHooks.contains("POST_TOOL"));
+    }
 }