Introduce hook and context management to OpenSearch Agents #4388
Conversation
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
requested review from
HenryL27,
Zhangxunmt,
austintlee,
b4sjoo,
dhrubo-os,
jngz-es,
model-collapse,
pyek-bot,
rbhavna,
sam-herman,
xinyual,
ylwu-amzn and
zane-neo
as code owners
mingshl
force-pushed
the
feature/context_manager_hooks_inline_create
branch
from
caa6d57 to
4eaeb94
Compare
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
Signed-off-by: Xun Zhang <xunzh@amazon.com>
…project#4345) * initiate context management api with hook implementation Signed-off-by: Mingshi Liu <mingshl@amazon.com> * apply spotless Signed-off-by: Mingshi Liu <mingshl@amazon.com> --------- Signed-off-by: Mingshi Liu <mingshl@amazon.com>
* add pre_llm hook to per agent Signed-off-by: Mingshi Liu <mingshl@amazon.com> change context management passing from query parameters to payload Signed-off-by: Mingshi Liu <mingshl@amazon.com> pass hook registery into PER Signed-off-by: Mingshi Liu <mingshl@amazon.com> apply spotless Signed-off-by: Mingshi Liu <mingshl@amazon.com> initiate context management api with hook implementation Signed-off-by: Mingshi Liu <mingshl@amazon.com> * add comment Signed-off-by: Mingshi Liu <mingshl@amazon.com> * format Signed-off-by: Mingshi Liu <mingshl@amazon.com> * add validation Signed-off-by: Mingshi Liu <mingshl@amazon.com> --------- Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
mingshl
force-pushed
the
feature/context_manager_hooks_inline_create
branch
from
2a1f919 to
491a1bb
Compare
WalkthroughAdds a pluggable context-management subsystem for ML agents: new context manager interfaces and implementations, activation rules, a hook framework, template CRUD APIs (transport + REST + index mappings), template service and factory, agent execution integration (HookRegistry wiring), validation, and tests. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client as NodeClient
participant MLExecute as MLExecuteTaskRunner
participant TemplateSvc as ContextManagementTemplateService
participant Factory as ContextManagerFactory
participant HookReg as HookRegistry
participant Agent as MLAgentRunner
Client->>MLExecute: execute task (AgentMLInput)
MLExecute->>MLExecute: determine effective context_management name
MLExecute->>TemplateSvc: getTemplate(name)
TemplateSvc-->>MLExecute: ContextManagementTemplate
MLExecute->>Factory: createContextManagers(template)
Factory-->>MLExecute: List<ContextManager>
MLExecute->>HookReg: create HookRegistry and register managers
MLExecute->>Agent: run agent with HookRegistry attached
Agent->>HookReg: emit PreLLMEvent(context)
HookReg->>ContextManager: shouldActivate / execute (per manager)
ContextManager-->>HookReg: may modify ContextManagerContext
Agent->>HookReg: emit PostToolEvent (EnhancedPostToolEvent) after tool
HookReg->>ContextManager: execute on post-tool
Agent-->>Client: final result
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Areas to focus during review:
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 14
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java (1)
139-174: Critical:hookRegistryis static, causing cross-request state sharing.The
hookRegistryfield is declared asprivate static HookRegistry hookRegistry. This means all instances ofMLChatAgentRunnershare the same registry. In a concurrent environment with multiple agent executions, this will cause:
- Cross-request contamination - hooks registered for one request affect others
- Race conditions when different requests try to use/modify the registry
- Memory leaks as hooks accumulate across requests
This should be an instance field.
- private static HookRegistry hookRegistry; + private HookRegistry hookRegistry; // ... constructors remain the same, but now each instance has its own hookRegistry
🟡 Minor comments (15)
common/src/main/java/org/opensearch/ml/common/input/execute/agent/AgentMLInput.java-112-120 (1)
112-120: Removecontext_managementfrom the parameters map after extracting it forcontextManagementName.The code extracts
context_managementintocontextManagementName(line 117) butStringUtils.getParameterMap()still includes it in the returnedparametersmap (line 114). Sincecontext_managementis a control parameter meant only for identifying the management template, it should be excluded from the actual inference parameters. Remove the key fromparameterObjsbefore passing it toStringUtils.getParameterMap(), or filter it out after conversion.ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java-396-428 (1)
396-428: The equality check will work correctly; formatting difference concern is incorrect.The assumption about formatting differences between
completedStepsandupdatedStepsis unfounded. Both lists contain the same<step-N>...</step-N>formatted strings when compared at line 412, sincecompletedStepsis passed directly toemitPreLLMHookand used as thetoolInteractionsreference inContextManagerContext. Theequals()check will properly detect if context managers modified the list.However, the TODO comment on line 420 indicating uncertainty about clearing the INTERACTIONS parameter should be resolved before merging, as it suggests the parameter handling logic remains incomplete.
ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/contextmanager/SlidingWindowManager.java-83-91 (1)
83-91: Duplicate empty check - second condition is unreachable.Lines 88-91 check
interactions.isEmpty()but this was already checked at line 83. If execution reaches line 88,interactionsis guaranteed to be non-null and non-empty, making lines 88-91 dead code.if (interactions == null || interactions.isEmpty()) { log.debug("No tool interactions to process"); return; } - - if (interactions.isEmpty()) { - log.debug("No string interactions found in tool interactions"); - return; - }common/src/main/java/org/opensearch/ml/common/contextmanager/ActivationRuleFactory.java-95-109 (1)
95-109:parseIntegerValuedoesn't handle null input.If
valueis null (e.g., key exists but value is null in config), line 107 will throwNullPointerExceptionwhen callingvalue.getClass(). Add explicit null handling:private static int parseIntegerValue(Object value, String fieldName) { + if (value == null) { + throw new IllegalArgumentException("Null value provided for " + fieldName); + } if (value instanceof Integer) {ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/contextmanager/SlidingWindowManager.java-126-147 (1)
126-147: Add null check forconfigparameter.If
configis null,config.get(key)at line 127 will throwNullPointerException. The similar method inToolsOutputTruncateManager(per relevant snippets) likely handles this case.private int parseIntegerConfig(Map<String, Object> config, String key, int defaultValue) { + if (config == null) { + return defaultValue; + } Object value = config.get(key);plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ListContextManagementTemplatesTransportAction.java-52-53 (1)
52-53: Pagination total count is incorrect.Using
templates.size()as the total means clients cannot determine the actual number of templates for pagination. Consider usingSearchResponse.getHits().getTotalHits().valuein the service layer to return the true total.plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ListContextManagementTemplatesTransportAction.java-47-61 (1)
47-61: Thread context not properly restored for async callbacks.The
try-with-resourcesblock closes immediately afterlistTemplatesreturns (before the async callbacks execute), so the context is restored prematurely. SinceContextManagementTemplateService.listTemplatesalready handles context stashing internally, this outer stash is redundant.@Override protected void doExecute( Task task, MLListContextManagementTemplatesRequest request, ActionListener<MLListContextManagementTemplatesResponse> listener ) { - try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + try { log.debug("Listing context management templates from: {} size: {}", request.getFrom(), request.getSize()); contextManagementTemplateService.listTemplates(request.getFrom(), request.getSize(), ActionListener.wrap(templates -> { log.debug("Successfully retrieved {} context management templates", templates.size()); - // For now, return the size as total. In a real implementation, you'd get the actual total count listener.onResponse(new MLListContextManagementTemplatesResponse(templates, templates.size())); }, exception -> { log.error("Error listing context management templates", exception); listener.onFailure(exception); })); } catch (Exception e) { log.error("Unexpected error listing context management templates", e); listener.onFailure(e); } }common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagementTemplate.java-246-255 (1)
246-255: Potential NPE: hooks map values could contain null entries.The validation iterates over
configsbut doesn't guard againstnullentries in the list. If a list contains anullconfig, callingconfig.isValid()will throwNullPointerException.// Validate all context manager configs for (List<ContextManagerConfig> configs : hooks.values()) { if (configs != null) { for (ContextManagerConfig config : configs) { + if (config == null) { + return false; + } if (!config.isValid()) { return false; } } } }common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java-154-173 (1)
154-173:updateHookConfigurationis not thread-safe and loses existing managers.This method clears
hookToManagersMapbefore repopulating it. During this window, any hook events would find no managers. Additionally, it callsfindManagerByTypewhich searches the originalcontextManagerslist, but won't find managers that weren't originally registered.Consider building a new map and replacing atomically:
public void updateHookConfiguration(Map<String, List<ContextManagerConfig>> hookConfiguration) { - hookToManagersMap.clear(); + Map<String, List<ContextManager>> newMap = new HashMap<>(); for (Map.Entry<String, List<ContextManagerConfig>> entry : hookConfiguration.entrySet()) { String hookName = entry.getKey(); List<ContextManagerConfig> configs = entry.getValue(); for (ContextManagerConfig config : configs) { ContextManager manager = findManagerByType(config.getType()); if (manager != null) { - addManagerToHook(hookName, manager); + newMap.computeIfAbsent(hookName, k -> new ArrayList<>()).add(manager); } else { log.warn("Context manager of type {} not found", config.getType()); } } } + hookToManagersMap.clear(); + hookToManagersMap.putAll(newMap); log.info("Updated hook configuration with {} hooks", hookConfiguration.size()); }plugin/src/main/java/org/opensearch/ml/task/MLExecuteTaskRunner.java-311-318 (1)
311-318: Passing null for NamedXContentRegistry may cause issues.Line 314 passes
nullas the first argument tocreateParser. While this may work for simple JSON parsing, it could causeNullPointerExceptionif the parser encounters named objects that require registry lookup.-org.opensearch.core.xcontent.XContentParser parser = - org.opensearch.common.xcontent.json.JsonXContent.jsonXContent - .createParser( - null, - org.opensearch.common.xcontent.LoggingDeprecationHandler.INSTANCE, - response.getSourceAsString() - ); +org.opensearch.core.xcontent.XContentParser parser = + org.opensearch.common.xcontent.json.JsonXContent.jsonXContent + .createParser( + org.opensearch.core.xcontent.NamedXContentRegistry.EMPTY, + org.opensearch.common.xcontent.LoggingDeprecationHandler.INSTANCE, + response.getSourceAsString() + );plugin/src/main/java/org/opensearch/ml/task/MLExecuteTaskRunner.java-348-350 (1)
348-350: Silent exception swallowing hides potential issues.The empty catch block silently swallows exceptions, making debugging difficult. At minimum, add debug logging to capture failure reasons.
-} catch (Exception e) { - // Continue to fallback methods -} +} catch (Exception e) { + log.debug("Failed to retrieve agent context management configuration for agentId: {}, falling back to other methods", agentId, e); +}ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/contextmanager/SummarizationManager.java-120-131 (1)
120-131: Redundant isEmpty() check - dead code.The check at line 128 is unreachable. If
interactions.isEmpty()is true at line 124, the method returns at line 125. The second check will never execute with an empty list.@Override public void execute(ContextManagerContext context) { List<String> interactions = context.getToolInteractions(); if (interactions == null || interactions.isEmpty()) { return; } - if (interactions.isEmpty()) { - log.debug("No string interactions found in tool interactions"); - return; - } - int totalMessages = interactions.size();ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java-168-169 (1)
168-169: Empty conditional block serves no purpose.This empty
ifblock forchatHistoryshould either be implemented or removed. Consider adding a TODO comment if this is intentional placeholder for future work.- if (context.getChatHistory() != null && !context.getChatHistory().isEmpty()) { - } + // TODO: Implement chat history propagation when context managers support itplugin/src/main/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateService.java-322-322 (1)
322-322: Unused variable:useris fetched but never used.The
uservariable fromRestActionUtils.getUserContext(client)is retrieved but not utilized. Either remove it or use it to setlastModifiedByon the template (similar tosaveTemplatesettingcreatedBy).public void updateTemplate(String templateName, ContextManagementTemplate template, ActionListener<UpdateResponse> listener) { try { // Validate template if (!template.isValid()) { listener.onFailure(new IllegalArgumentException("Invalid context management template")); return; } - User user = RestActionUtils.getUserContext(client); try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {common/src/test/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentRequestTest.java-297-318 (1)
297-318: Remove duplicate test.The tests
validate_ContextManagementTemplateName_NullValue(lines 297-306) andvalidate_ContextManagementTemplateName_Null(lines 308-318) are identical duplicates that test the same scenario with the same assertions.Remove one of the duplicate tests:
- @Test - public void validate_ContextManagementTemplateName_Null() { - // Test null template name validation - MLAgent agentWithNullName = MLAgent.builder().name("test_agent").type("flow").contextManagementName(null).build(); - - MLRegisterAgentRequest request = new MLRegisterAgentRequest(agentWithNullName); - ActionRequestValidationException exception = request.validate(); - - // This should pass since null is handled differently than empty - assertNull(exception); - } -
common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerContext.java
Show resolved
Hide resolved
common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java
Show resolved
Hide resolved
common/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentRequest.java
Show resolved
Hide resolved
ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java
Show resolved
Hide resolved
ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java
Show resolved
Hide resolved
...g/opensearch/ml/action/contextmanagement/DeleteContextManagementTemplateTransportAction.java
Show resolved
Hide resolved
.../org/opensearch/ml/action/contextmanagement/GetContextManagementTemplateTransportAction.java
Outdated
Show resolved
Hide resolved
plugin/src/main/java/org/opensearch/ml/rest/RestMLListContextManagementTemplatesAction.java
Show resolved
Hide resolved
...t/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateServiceTests.java
Show resolved
Hide resolved
Signed-off-by: Mingshi Liu <mingshl@amazon.com>
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Error
mingshl
had a problem deploying
to
ml-commons-cicd-env
GitHub Actions
Failure
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (2)
common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java (2)
154-173: Non-atomic update inupdateHookConfigurationcan cause inconsistent reads.The
clear()followed by iteration allows readers (hook event handlers) to see an empty or partially populated map during the update. Consider building a new map and atomically swapping:public void updateHookConfiguration(Map<String, List<ContextManagerConfig>> hookConfiguration) { - hookToManagersMap.clear(); + Map<String, List<ContextManager>> newMap = new ConcurrentHashMap<>(); for (Map.Entry<String, List<ContextManagerConfig>> entry : hookConfiguration.entrySet()) { String hookName = entry.getKey(); List<ContextManagerConfig> configs = entry.getValue(); for (ContextManagerConfig config : configs) { ContextManager manager = findManagerByType(config.getType()); if (manager != null) { - addManagerToHook(hookName, manager); + newMap.computeIfAbsent(hookName, k -> new CopyOnWriteArrayList<>()).add(manager); } else { log.warn("Context manager of type {} not found", config.getType()); } } } + // Atomically replace all mappings + hookToManagersMap.clear(); + hookToManagersMap.putAll(newMap); + log.info("Updated hook configuration with {} hooks", hookConfiguration.size()); }Note: For true atomicity, consider using
AtomicReference<Map<String, List<ContextManager>>>and swapping the entire reference.
145-148: Thread safety:ArrayListvalues inConcurrentHashMapare not thread-safe.While
hookToManagersMapusesConcurrentHashMap, theArrayListvalues created incomputeIfAbsentare not thread-safe. IfaddManagerToHookis called concurrently, or ifexecuteManagersForHookiterates while the list is being modified, aConcurrentModificationExceptioncould occur.Consider using
CopyOnWriteArrayList:private void addManagerToHook(String hookName, ContextManager manager) { - hookToManagersMap.computeIfAbsent(hookName, k -> new ArrayList<>()).add(manager); + hookToManagersMap.computeIfAbsent(hookName, k -> new CopyOnWriteArrayList<>()).add(manager); log.debug("Added manager {} to hook {}", manager.getType(), hookName); }
🧹 Nitpick comments (11)
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java (1)
745-745: Consider verifying specific parameters instead of using magic numbers.These four tests now assert specific parameter map sizes (17, 18, 18, 18) to account for the additional context-management parameter. However, verifying only the count is brittle and provides limited coverage of the new feature.
Consider refactoring to:
- Define expected parameter keys as constants or verify their presence explicitly
- Assert that the context-management parameter exists and has the expected value
- Add a comment explaining why different tests expect different counts
Example for line 745:
- assertEquals(17, ((Map) argumentCaptor.getValue()).size()); + Map<String, String> toolParams = (Map<String, String>) argumentCaptor.getValue(); + // Verify context management parameter is injected + assertTrue("Context management parameter should be present", toolParams.containsKey("expected_context_key")); + // Verify other expected parameters + assertTrue(toolParams.containsKey("key1")); + assertTrue(toolParams.containsKey("key2"));This approach makes tests more maintainable and provides better documentation of the parameter contract.
Also applies to: 773-773, 839-839, 869-869
common/src/test/java/org/opensearch/ml/common/contextmanager/ContextManagerContextTest.java (2)
39-64: Consider adding test forgetMessageCountwith builder pattern.There's a test for
getEstimatedTokenCountwith both no-arg constructor and builder, butgetMessageCountonly has a test for the no-arg constructor. For consistency and completeness, consider adding a similar test:@Test public void testGetMessageCountWithBuilder() { ContextManagerContext context = ContextManagerContext.builder().build(); // Should not throw NPE int messageCount = context.getMessageCount(); assertEquals(0, messageCount); }
7-65: Test coverage is good for null-safety, but consider adding tests with populated data.The tests verify that collections are properly initialized and methods don't throw NPE on empty contexts. Consider adding tests that populate the context with actual chat history and tool interactions to verify the token and message counting logic works correctly with real data.
plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateService.java (2)
321-321: Unused variable:useris retrieved but never used.The
uservariable is fetched from the client context but isn't used in theupdateTemplatemethod. Either remove it or use it (e.g., to track who modified the template).- User user = RestActionUtils.getUserContext(client);
285-305: Consider usingexistsAPI instead of search for uniqueness check.Since templates are stored with
template.getName()as the document_id(line 110), you could use a simplerGetRequestwithfetchSourceContext(FetchSourceContext.DO_NOT_FETCH_SOURCE)or the exists API instead of a search query. This would be more efficient:GetRequest getRequest = new GetRequest(ContextManagementIndexUtils.getIndexName(), templateName) .fetchSourceContext(FetchSourceContext.DO_NOT_FETCH_SOURCE); client.get(getRequest, ActionListener.wrap( response -> listener.onResponse(response.isExists()), // ... error handling ));plugin/src/test/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateServiceTests.java (1)
271-281: Test doesn't assert expected behavior for null listener.This test catches any exception without asserting specific behavior. If null listeners should be rejected, the test should verify an
IllegalArgumentExceptionorNullPointerExceptionis thrown. If they should be gracefully handled, the test should verify no exception occurs.@Test public void testListTemplates_NullListener() { // Either assert it throws a specific exception: assertThrows(NullPointerException.class, () -> contextManagementTemplateService.listTemplates(null)); // Or verify it handles gracefully without throwing }ml-algorithms/src/test/java/org/opensearch/ml/engine/agents/AgentContextUtilTest.java (2)
34-49: Consider verifying the event type passed toemit().The test verifies that
emit()is called once but doesn't verify that aPreLLMEventis passed. This would strengthen the test:- verify(hookRegistry, times(1)).emit(any()); + ArgumentCaptor<PreLLMEvent> eventCaptor = ArgumentCaptor.forClass(PreLLMEvent.class); + verify(hookRegistry, times(1)).emit(eventCaptor.capture()); + assertNotNull(eventCaptor.getValue()); + assertEquals("test question", eventCaptor.getValue().getContext().getUserPrompt());
17-49: Consider adding test for exception handling path.Based on the implementation in
AgentContextUtil.emitPreLLMHook, if an exception occurs during hook emission, it logs the error and returns the original context. Consider adding a test for this path:@Test public void testEmitPreLLMHookWithHookRegistryException() { Map<String, String> parameters = new HashMap<>(); parameters.put("question", "test question"); HookRegistry hookRegistry = mock(HookRegistry.class); doThrow(new RuntimeException("Hook error")).when(hookRegistry).emit(any()); ContextManagerContext result = AgentContextUtil.emitPreLLMHook( parameters, new ArrayList<>(), new ArrayList<>(), mock(Memory.class), hookRegistry); assertNotNull(result); assertEquals("test question", result.getUserPrompt()); }plugin/src/main/java/org/opensearch/ml/action/contextmanagement/UpdateContextManagementTemplateTransportAction.java (1)
42-54: Consider using try-with-resources for consistent thread context handling.The current pattern manually stashes and restores the context. If
updateTemplatethrows synchronously before initiating its async operation, and the exception is caught at line 49, the context is restored at line 51. However, using try-with-resources would be more consistent with the pattern used inContextManagementTemplateService:@Override protected void doExecute(Task task, MLUpdateContextManagementTemplateRequest request, ActionListener<UpdateResponse> listener) { - ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext(); - ActionListener<UpdateResponse> wrappedListener = ActionListener.runBefore(listener, context::restore); - try { + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + ActionListener<UpdateResponse> wrappedListener = ActionListener.runBefore(listener, context::restore); log.info("Updating context management template: {}", request.getTemplateName()); - contextManagementTemplateService.updateTemplate(request.getTemplateName(), request.getTemplate(), wrappedListener); } catch (Exception e) { log.error("Failed to update context management template: {}", request.getTemplateName(), e); - context.restore(); listener.onFailure(e); } }ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java (2)
27-55: Unusedmemoryparameter.The
memoryparameter is declared but never referenced in the method body. If it's reserved for future use, consider adding a comment or removing it to reduce confusion.
172-173: Remove empty if block.This empty conditional serves no purpose and should be removed to improve code clarity.
Apply this diff:
- if (context.getChatHistory() != null && !context.getChatHistory().isEmpty()) { - } -
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java(1 hunks)common/src/test/java/org/opensearch/ml/common/contextmanager/ContextManagerContextTest.java(1 hunks)ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java(1 hunks)ml-algorithms/src/test/java/org/opensearch/ml/engine/agents/AgentContextUtilTest.java(1 hunks)ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java(4 hunks)plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateService.java(1 hunks)plugin/src/main/java/org/opensearch/ml/action/contextmanagement/GetContextManagementTemplateTransportAction.java(1 hunks)plugin/src/main/java/org/opensearch/ml/action/contextmanagement/UpdateContextManagementTemplateTransportAction.java(1 hunks)plugin/src/main/java/org/opensearch/ml/rest/RestMLListContextManagementTemplatesAction.java(1 hunks)plugin/src/test/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateServiceTests.java(1 hunks)plugin/src/test/java/org/opensearch/ml/rest/RestMLListContextManagementTemplatesActionTests.java(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- plugin/src/test/java/org/opensearch/ml/rest/RestMLListContextManagementTemplatesActionTests.java
- plugin/src/main/java/org/opensearch/ml/rest/RestMLListContextManagementTemplatesAction.java
- plugin/src/main/java/org/opensearch/ml/action/contextmanagement/GetContextManagementTemplateTransportAction.java
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-12-15T23:04:24.444Z
Learnt from: mingshl
Repo: opensearch-project/ml-commons PR: 4388
File: common/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentRequest.java:102-105
Timestamp: 2025-12-15T23:04:24.444Z
Learning: In MLRegisterAgentRequest.isValidHookName(), POST_MEMORY is intentionally excluded from validation because this hook is not yet used in context managers, even though infrastructure exists for it in ContextManagerHookProvider.
Applied to files:
ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java
🧬 Code graph analysis (3)
ml-algorithms/src/test/java/org/opensearch/ml/engine/agents/AgentContextUtilTest.java (1)
ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java (1)
AgentContextUtil(24-186)
common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java (10)
common/src/main/java/org/opensearch/ml/common/hooks/EnhancedPostToolEvent.java (1)
EnhancedPostToolEvent(19-46)common/src/main/java/org/opensearch/ml/common/hooks/PostMemoryEvent.java (1)
PostMemoryEvent(19-50)common/src/main/java/org/opensearch/ml/common/hooks/PreLLMEvent.java (1)
PreLLMEvent(17-37)plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateService.java (1)
Log4j2(48-358)ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/contextmanager/SlidingWindowManager.java (1)
Log4j2(25-191)ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/contextmanager/SummarizationManager.java (1)
Log4j2(49-435)ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/contextmanager/ToolsOutputTruncateManager.java (1)
Log4j2(22-134)common/src/main/java/org/opensearch/ml/common/hooks/HookRegistry.java (1)
Log4j2(20-93)plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ContextManagementIndexUtils.java (1)
Log4j2(21-96)plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ContextManagerFactory.java (1)
Log4j2(26-120)
ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java (3)
common/src/main/java/org/opensearch/ml/common/input/execute/agent/AgentMLInput.java (1)
org(31-131)common/src/main/java/org/opensearch/ml/common/hooks/EnhancedPostToolEvent.java (1)
EnhancedPostToolEvent(19-46)common/src/main/java/org/opensearch/ml/common/hooks/PreLLMEvent.java (1)
PreLLMEvent(17-37)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: spotless
🔇 Additional comments (3)
common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java (1)
90-112: Good fault-tolerant execution pattern.The
executeManagersForHookmethod properly isolates failures - if one context manager fails, it logs the error and continues with the remaining managers. This aligns well with the PR objective of fault-tolerant hook handling.ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java (2)
139-161: Null check properly implemented.The null check for
hookRegistryhas been correctly added, addressing the previous review concern. The method now safely handles null registries.
175-177: No changes needed. The leading comma in theINTERACTIONSparameter is intentional. This parameter is used in template variable substitution viaINTERACTIONS_PREFIX("${_interactions.") where the comma prefix serves as a separator for formatting interactions that will be embedded into prompt templates. The value is meant to be processed by the context management system and integrated into prompts, making the leading comma a necessary part of the designed format.Likely an incorrect or invalid review comment.
| } else if ("SlidingWindowManager".equals(managerType) || "SummarizingManager".equals(managerType)) { | ||
| addManagerToHook("POST_MEMORY", manager); | ||
| addManagerToHook("PRE_LLM", manager); |
There was a problem hiding this comment.
Typo: "SummarizingManager" should be "SummarizationManager".
Based on the SummarizationManager class which defines TYPE = "SummarizationManager", the comparison here has a typo:
- } else if ("SlidingWindowManager".equals(managerType) || "SummarizingManager".equals(managerType)) {
+ } else if ("SlidingWindowManager".equals(managerType) || "SummarizationManager".equals(managerType)) {This would cause SummarizationManager instances to fall through to the default case (line 135) and only be registered for PRE_LLM, missing the intended POST_MEMORY hook.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| } else if ("SlidingWindowManager".equals(managerType) || "SummarizingManager".equals(managerType)) { | |
| addManagerToHook("POST_MEMORY", manager); | |
| addManagerToHook("PRE_LLM", manager); | |
| } else if ("SlidingWindowManager".equals(managerType) || "SummarizationManager".equals(managerType)) { | |
| addManagerToHook("POST_MEMORY", manager); | |
| addManagerToHook("PRE_LLM", manager); |
🤖 Prompt for AI Agents
In
common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerHookProvider.java
around lines 128 to 130, the string literal "SummarizingManager" is a typo and
should be "SummarizationManager"; update the equality check to use
"SummarizationManager" so instances of SummarizationManager match that branch
and get registered for both POST_MEMORY and PRE_LLM hooks.
| public void saveTemplate(String templateName, ContextManagementTemplate template, ActionListener<Boolean> listener) { | ||
| try { | ||
| // Validate template | ||
| if (!template.isValid()) { | ||
| listener.onFailure(new IllegalArgumentException("Invalid context management template")); | ||
| return; | ||
| } |
There was a problem hiding this comment.
Missing validation for templateName parameter in saveTemplate.
The method validates template.isValid() but doesn't validate the templateName parameter for null/empty/whitespace like getTemplate and deleteTemplate do. Additionally, there's potential for inconsistency if templateName differs from template.getName().
Consider adding validation:
public void saveTemplate(String templateName, ContextManagementTemplate template, ActionListener<Boolean> listener) {
try {
+ if (templateName == null || templateName.trim().isEmpty()) {
+ listener.onFailure(new IllegalArgumentException("Template name cannot be null, empty, or whitespace"));
+ return;
+ }
+
// Validate template
if (!template.isValid()) {
listener.onFailure(new IllegalArgumentException("Invalid context management template"));
return;
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| public void saveTemplate(String templateName, ContextManagementTemplate template, ActionListener<Boolean> listener) { | |
| try { | |
| // Validate template | |
| if (!template.isValid()) { | |
| listener.onFailure(new IllegalArgumentException("Invalid context management template")); | |
| return; | |
| } | |
| public void saveTemplate(String templateName, ContextManagementTemplate template, ActionListener<Boolean> listener) { | |
| try { | |
| if (templateName == null || templateName.trim().isEmpty()) { | |
| listener.onFailure(new IllegalArgumentException("Template name cannot be null, empty, or whitespace")); | |
| return; | |
| } | |
| // Validate template | |
| if (!template.isValid()) { | |
| listener.onFailure(new IllegalArgumentException("Invalid context management template")); | |
| return; | |
| } |
🤖 Prompt for AI Agents
In
plugin/src/main/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateService.java
around lines 70 to 76, add validation for the templateName parameter (null,
empty, or only whitespace) similar to getTemplate/deleteTemplate and ensure
consistency between templateName and template.getName(): if templateName is
invalid call listener.onFailure(new IllegalArgumentException(...)); if
template.getName() is null or empty, set it from templateName or if both are
present and differ, fail with an IllegalArgumentException explaining the
mismatch; proceed only after these checks and use the validated name for saving.
| public void updateTemplate(String templateName, ContextManagementTemplate template, ActionListener<UpdateResponse> listener) { | ||
| try { | ||
| // Validate template | ||
| if (!template.isValid()) { | ||
| listener.onFailure(new IllegalArgumentException("Invalid context management template")); | ||
| return; | ||
| } | ||
|
|
||
| User user = RestActionUtils.getUserContext(client); | ||
| try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { | ||
| ActionListener<UpdateResponse> wrappedListener = ActionListener.runBefore(listener, context::restore); | ||
|
|
||
| // Set last modified timestamp | ||
| template.setLastModified(Instant.now()); | ||
|
|
||
| // Create the update request | ||
| UpdateRequest updateRequest = new UpdateRequest(ContextManagementIndexUtils.getIndexName(), templateName) | ||
| .doc(template.toXContent(jsonXContent.contentBuilder(), ToXContentObject.EMPTY_PARAMS)) | ||
| .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE); | ||
|
|
||
| // Execute the update operation | ||
| client.update(updateRequest, ActionListener.wrap(updateResponse -> { | ||
| log.info("Context management template updated successfully: {}", templateName); | ||
| wrappedListener.onResponse(updateResponse); | ||
| }, exception -> { | ||
| log.error("Failed to update context management template: {}", templateName, exception); | ||
| wrappedListener.onFailure(exception); | ||
| })); | ||
| } | ||
| } catch (Exception e) { | ||
| log.error("Error updating context management template: {}", templateName, e); | ||
| listener.onFailure(e); | ||
| } | ||
| } |
There was a problem hiding this comment.
updateTemplate lacks input validation and explicit "not found" handling.
Similar to saveTemplate, the templateName parameter isn't validated. Additionally, unlike deleteTemplate which handles IndexNotFoundException, this method doesn't handle the case where the template doesn't exist (will throw DocumentMissingException).
Consider adding validation and consistent error handling:
public void updateTemplate(String templateName, ContextManagementTemplate template, ActionListener<UpdateResponse> listener) {
try {
+ if (templateName == null || templateName.trim().isEmpty()) {
+ listener.onFailure(new IllegalArgumentException("Template name cannot be null, empty, or whitespace"));
+ return;
+ }
+
// Validate template
if (!template.isValid()) {And handle document missing exception in the error callback for better UX.
| @Test | ||
| public void testGetTemplate_EmptyTemplateName() { | ||
| @SuppressWarnings("unchecked") | ||
| ActionListener<ContextManagementTemplate> listener = mock(ActionListener.class); | ||
|
|
||
| contextManagementTemplateService.getTemplate("", listener); | ||
|
|
||
| ArgumentCaptor<IllegalArgumentException> exceptionCaptor = ArgumentCaptor.forClass(IllegalArgumentException.class); | ||
| verify(listener).onFailure(exceptionCaptor.capture()); | ||
| assertEquals("Template name cannot be null, empty, or whitespace", exceptionCaptor.getValue().getMessage()); | ||
| } |
There was a problem hiding this comment.
Duplicate test methods for empty template name validation.
testGetTemplate_EmptyTemplateName is defined twice (lines 171-181 and 296-304). The second occurrence should be removed or renamed if testing different behavior.
- @Test
- public void testGetTemplate_EmptyTemplateName() {
- @SuppressWarnings("unchecked")
- ActionListener<ContextManagementTemplate> listener = mock(ActionListener.class);
-
- contextManagementTemplateService.getTemplate("", listener);
-
- verify(listener).onFailure(any(IllegalArgumentException.class));
- verify(client, never()).threadPool();
- }Also applies to: 296-304
🤖 Prompt for AI Agents
In
plugin/src/test/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateServiceTests.java
around lines 171-181 and 296-304, there is a duplicate test method named
testGetTemplate_EmptyTemplateName; remove the redundant second definition (lines
296-304) or if it was intended to cover a different case, rename it and adjust
its assertions to reflect the distinct behavior being tested, then run the test
suite to ensure no name collisions or assertion mismatches remain.
| @Test | ||
| public void testDeleteTemplate_EmptyTemplateName() { | ||
| @SuppressWarnings("unchecked") | ||
| ActionListener<Boolean> listener = mock(ActionListener.class); | ||
|
|
||
| contextManagementTemplateService.deleteTemplate("", listener); | ||
|
|
||
| ArgumentCaptor<IllegalArgumentException> exceptionCaptor = ArgumentCaptor.forClass(IllegalArgumentException.class); | ||
| verify(listener).onFailure(exceptionCaptor.capture()); | ||
| assertEquals("Template name cannot be null, empty, or whitespace", exceptionCaptor.getValue().getMessage()); | ||
| } |
There was a problem hiding this comment.
Duplicate test methods for delete with empty template name.
testDeleteTemplate_EmptyTemplateName is also defined twice (lines 195-205 and 318-326). Remove the duplicate.
Also applies to: 318-326
🤖 Prompt for AI Agents
In
plugin/src/test/java/org/opensearch/ml/action/contextmanagement/ContextManagementTemplateServiceTests.java
around lines 195-205 (and duplicated at 318-326), there are two identical test
methods named testDeleteTemplate_EmptyTemplateName; remove one of the duplicates
(delete the redundant method block at either 195-205 or 318-326) so the test
only appears once, and ensure any related imports or test annotations remain
intact after removal.
3 participants
Description
Add Complete Context Management Framework with Memory-Integrated Hook Events to ML Agents
This PR introduces a comprehensive context management system for ML agents, featuring memory-integrated hook
events, complete hook registry with fault-tolerant execution, enhanced context integration, and intelligent
resource management for enterprise-grade deployments.
Major Features:
Complete Hook Registry System
• Fault-Tolerant Event Emission: emitEvent() with exception isolation and comprehensive logging
• Thread-Safe Operations: ConcurrentHashMap ensuring concurrent callback registration and execution
• Callback Management: Type-safe callback registration with monitoring capabilities
• Production-Grade Error Handling: Individual callback failures don't affect system operation
Comprehensive Hook Event Architecture
• HookEvent Base Class: Abstract base with invocation state for all lifecycle events
• PreLLMEvent: pre llm event with context integration and retrieved re-act interactions, including tool used result/ completed steps
• EnhancedPostToolEvent: Tool event with context integration
• Event Hierarchy: Well-designed inheritance from HookEvent to specialized events
Complete CRUD REST API Implementation
• Create Operations: MLCreateContextManagementTemplateAction, MLCreateContextManagementTemplateRequest,
MLCreateContextManagementTemplateResponse
• Read Operations: MLGetContextManagementTemplateAction, MLGetContextManagementTemplateRequest,
MLGetContextManagementTemplateResponse
• Delete Operations: MLDeleteContextManagementTemplateAction, MLDeleteContextManagementTemplateRequest,
MLDeleteContextManagementTemplateResponse
• Full CRUD Lifecycle: Complete template management with create, read, and delete operations
Context Management Integration
• ML Agent Integration: contextManagementName and contextManagement fields with validation
• Template System: ContextManagementTemplate with hooks, validation, and serialization
• Hook Provider Integration: ContextManagerHookProvider registers callbacks for all event types
Rich Context State Management
• ContextManagerContext: Comprehensive context with prompts, history, tools, interactions, and parameters
• Utility Methods: getEstimatedTokenCount(), getMessageCount(), addParameter(), addToolInteraction()
• Memory Integration: Context managers can access both current context and retrieved memory
• Dynamic Modification: Runtime context updates with memory and invocation state tracking
Event-Driven Execution Pipeline
• Tool Event Processing: handlePostTool() processes enhanced tool events with context
• Pre-LLM Processing: handlePreLLM() processes context before LLM execution
• Memory Event Processing: handlePostMemory() processes memory events with context and history
• Automatic Execution: All events processed automatically through hook registry
Complete Activation Rules System
• TokensExceedRule: Context-aware activation using comprehensive token estimation
• MessageCountExceedRule: Conversation length management with memory integration
• CompositeActivationRule: Complex conditions considering memory and context state
Token Counting Architecture
• TokenCounter Interface: Standardized interface with multiple truncation strategies
• Context Integration: Token counters integrated with activation rules and memory state
• Pluggable Design: Interface supports model-specific implementations
Pluggable Context Manager Architecture
• ContextManagerConfig: Configuration with type, activation, and manager settings
• Template Management: Complete template system with memory-integrated hooks
• Registry Integration: Managers registered for memory, tool, and LLM events
Technical Highlights:
• Comprehensive Event Model: Complete lifecycle coverage with memory, tool, and LLM events
• Fault-Tolerant Processing: Exception isolation ensures system stability across all event types
• Rich Context Access: Context managers have access to full context and memory
Architecture Benefits:
• Complete Lifecycle Coverage: Events cover all major phases of ML agent execution
• Integrated Memory Management: Memory retrieval and context management work together seamlessly
• Enterprise Reliability: Fault-tolerant design with comprehensive error handling
Use Cases:
• Comprehensive Resource Management: Manage tokens across current context and memory history
Impact:
• Comprehensive Agent Lifecycle: Full coverage of memory, tool, and LLM execution phases
• Advanced AI Orchestration: Sophisticated context management with memory integration
Check List
--signoff.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.
Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.