address comments

jiapingzeng · jiapingzeng · commit 3bddf55ce6cb · 2025-11-19T01:27:57.000-08:00
Signed-off-by: Jiaping Zeng &lt;jpz@amazon.com&gt;
diff --git a/common/src/main/java/org/opensearch/ml/common/agui/AGUIInputConverter.java b/common/src/main/java/org/opensearch/ml/common/agui/AGUIInputConverter.java
@@ -72,7 +72,7 @@ public static boolean isAGUIInput(String inputJson) {
 
             return true;
         } catch (Exception e) {
-            log.debug("Failed to parse input as JSON for AG-UI detection", e);
+            log.error("Failed to parse input as JSON for AG-UI detection", e);
             return false;
         }
     }
diff --git a/common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java b/common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java
@@ -7,6 +7,7 @@
 
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_MEMORY_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENT_FRAMEWORK_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_EXECUTE_TOOL_ENABLED;
@@ -22,7 +23,6 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_REMOTE_INFERENCE_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STREAM_ENABLED;
-import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_ENABLED;
 
 import java.util.ArrayList;
 import java.util.List;
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java
@@ -1066,6 +1066,8 @@ public static Map<String, Tool> wrapFrontendToolsAsToolObjects(List<Map<String,
         }
 
         return wrappedTools;
+    }
+
     public static Map<String, Object> createMemoryParams(
         String question,
         String memoryId,
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAGUIAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAGUIAgentRunner.java
@@ -37,7 +37,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.agent.MLAgent;
-import org.opensearch.ml.common.spi.memory.Memory;
+import org.opensearch.ml.common.memory.Memory;
 import org.opensearch.ml.common.spi.tools.Tool;
 import org.opensearch.ml.engine.encryptor.Encryptor;
 import org.opensearch.ml.engine.function_calling.FunctionCalling;
@@ -362,6 +362,7 @@ private void processAGUIMessages(Map<String, String> params, String llmInterface
             }
         } catch (Exception e) {
             log.error("Failed to process AG-UI messages to chat history", e);
+            throw new IllegalArgumentException("Failed to process AG-UI messages to chat history", e);
         }
     }
 
@@ -405,6 +406,7 @@ private void processAGUIContext(Map<String, String> params) {
 
         } catch (Exception e) {
             log.error("Failed to process AG-UI context", e);
+            throw new IllegalArgumentException("Failed to process AG-UI context", e);
         }
     }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteAction.java
@@ -8,10 +8,10 @@
 import static org.opensearch.core.rest.RestStatus.BAD_REQUEST;
 import static org.opensearch.core.rest.RestStatus.INTERNAL_SERVER_ERROR;
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_DISABLED_MESSAGE;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_EXECUTE_TOOL_DISABLED_MESSAGE;
 import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
 import static org.opensearch.ml.utils.MLExceptionUtils.AGENT_FRAMEWORK_DISABLED_ERR_MSG;
-import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_AGENT_ID;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_ALGORITHM;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_TOOL_NAME;
@@ -140,6 +140,13 @@ MLExecuteTaskRequest getRequest(RestRequest request) throws IOException {
                 );
             } else {
                 input = MLInput.parse(parser, functionName.name());
+
+                if (!(input instanceof AgentMLInput)) {
+                    throw new IllegalArgumentException(
+                        String.format("Invalid input type. Expected: AgentMLInput, Received: %s", input.getClass().getSimpleName())
+                    );
+                }
+
                 ((AgentMLInput) input).setAgentId(agentId);
                 ((AgentMLInput) input).setTenantId(tenantId);
                 ((AgentMLInput) input).setIsAsync(async);
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteStreamAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteStreamAction.java
@@ -12,11 +12,11 @@
 import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_BACKEND_TOOL_NAMES;
 import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_RUN_ID;
 import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_THREAD_ID;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_DISABLED_MESSAGE;
 import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
 import static org.opensearch.ml.plugin.MachineLearningPlugin.STREAM_EXECUTE_THREAD_POOL;
 import static org.opensearch.ml.utils.MLExceptionUtils.AGENT_FRAMEWORK_DISABLED_ERR_MSG;
 import static org.opensearch.ml.utils.MLExceptionUtils.STREAM_DISABLED_ERR_MSG;
-import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AG_UI_DISABLED_MESSAGE;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_AGENT_ID;
 import static org.opensearch.ml.utils.RestActionUtils.isAsync;
 import static org.opensearch.ml.utils.TenantAwareHelper.getTenantID;
@@ -182,22 +182,41 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
                         BaseEvent runStartedEvent = new RunStartedEvent(threadId, runId);
                         HttpChunk startChunk = createHttpChunk("data: " + runStartedEvent.toJsonString() + "\n\n", false);
                         channel.sendChunk(startChunk);
-                        log.debug("RestMLExecuteStreamAction: Sent RUN_STARTED event - threadId={}, runId={}", threadId, runId);
+                        log.debug("AG-UI: RestMLExecuteStreamAction: Sent RUN_STARTED event - threadId={}, runId={}", threadId, runId);
                     }
 
                     // Extract backend tool names from agent configuration and add to request for AG-UI filtering
                     List<String> backendToolNames = extractBackendToolNamesFromAgent(agent);
                     if (isAGUI && !backendToolNames.isEmpty()) {
                         // Add backend tool names to request parameters so they're available during streaming
-                        RemoteInferenceInputDataSet inputDataSet = (RemoteInferenceInputDataSet) ((AgentMLInput) mlExecuteTaskRequest
-                            .getInput()).getInputDataset();
-                        inputDataSet.getParameters().put(AGUI_PARAM_BACKEND_TOOL_NAMES, new Gson().toJson(backendToolNames));
-                        log
-                            .info(
-                                "AG-UI: Added {} backend tool names to request for streaming filter: {}",
-                                backendToolNames.size(),
-                                backendToolNames
-                            );
+                        try {
+                            if (!(mlExecuteTaskRequest.getInput() instanceof AgentMLInput)) {
+                                throw new IllegalArgumentException(
+                                    "Invalid input type. Expected: AgentMLInput, Received: "
+                                        + mlExecuteTaskRequest.getInput().getClass().getSimpleName()
+                                );
+                            }
+                            AgentMLInput agentInput = (AgentMLInput) mlExecuteTaskRequest.getInput();
+
+                            if (!(agentInput.getInputDataset() instanceof RemoteInferenceInputDataSet)) {
+                                throw new IllegalArgumentException(
+                                    "Invalid dataset type. Expected: RemoteInferenceInputDataSet, Received: "
+                                        + agentInput.getInputDataset().getClass().getSimpleName()
+                                );
+                            }
+                            RemoteInferenceInputDataSet inputDataSet = (RemoteInferenceInputDataSet) agentInput.getInputDataset();
+
+                            inputDataSet.getParameters().put(AGUI_PARAM_BACKEND_TOOL_NAMES, new Gson().toJson(backendToolNames));
+                            log
+                                .info(
+                                    "AG-UI: Added {} backend tool names to request for streaming filter: {}",
+                                    backendToolNames.size(),
+                                    backendToolNames
+                                );
+                        } catch (ClassCastException e) {
+                            log.error("Failed to cast input types for backend tool names extraction", e);
+                            throw new IllegalArgumentException("Invalid input type configuration for AG-UI request", e);
+                        }
                     }
 
                     final CompletableFuture<HttpChunk> future = new CompletableFuture<>();

Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@ public static boolean isAGUIInput(String inputJson) {`
`72`	`72`
`73`	`73`	`return true;`
`74`	`74`	`} catch (Exception e) {`
`75`		`- log.debug("Failed to parse input as JSON for AG-UI detection", e);`
	`75`	`+ log.error("Failed to parse input as JSON for AG-UI detection", e);`
`76`	`76`	`return false;`
`77`	`77`	`}`
`78`	`78`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1066,6 +1066,8 @@ public static Map<String, Tool> wrapFrontendToolsAsToolObjects(List<Map<String,`
`1066`	`1066`	`}`
`1067`	`1067`
`1068`	`1068`	`return wrappedTools;`
	`1069`	`+ }`
	`1070`	`+`
`1069`	`1071`	`public static Map<String, Object> createMemoryParams(`
`1070`	`1072`	`String question,`
`1071`	`1073`	`String memoryId,`
Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@`
`37`	`37`	`import org.opensearch.core.action.ActionListener;`
`38`	`38`	`import org.opensearch.core.xcontent.NamedXContentRegistry;`
`39`	`39`	`import org.opensearch.ml.common.agent.MLAgent;`
`40`		`-import org.opensearch.ml.common.spi.memory.Memory;`
	`40`	`+import org.opensearch.ml.common.memory.Memory;`
`41`	`41`	`import org.opensearch.ml.common.spi.tools.Tool;`
`42`	`42`	`import org.opensearch.ml.engine.encryptor.Encryptor;`
`43`	`43`	`import org.opensearch.ml.engine.function_calling.FunctionCalling;`
`@@ -362,6 +362,7 @@ private void processAGUIMessages(Map<String, String> params, String llmInterface`
`362`	`362`	`}`
`363`	`363`	`} catch (Exception e) {`
`364`	`364`	`log.error("Failed to process AG-UI messages to chat history", e);`
	`365`	`+ throw new IllegalArgumentException("Failed to process AG-UI messages to chat history", e);`
`365`	`366`	`}`
`366`	`367`	`}`
`367`	`368`
`@@ -405,6 +406,7 @@ private void processAGUIContext(Map<String, String> params) {`
`405`	`406`
`406`	`407`	`} catch (Exception e) {`
`407`	`408`	`log.error("Failed to process AG-UI context", e);`
	`409`	`+ throw new IllegalArgumentException("Failed to process AG-UI context", e);`
`408`	`410`	`}`
`409`	`411`	`}`
`410`	`412`	`}`