Sanitize error messages to avoid reflecting user input

Author: JishaMNair

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

@@ -12,13 +12,7 @@
 
 import java.io.IOException;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Locale;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
+import java.util.*;
 
 import org.opensearch.Version;
 import org.opensearch.core.common.io.stream.StreamInput;
@@ -122,7 +116,7 @@ private void validate() {
             for (MLToolSpec toolSpec : tools) {
                 String toolName = Optional.ofNullable(toolSpec.getName()).orElse(toolSpec.getType());
                 if (toolNames.contains(toolName)) {
-                    throw new IllegalArgumentException("Duplicate tool defined: " + toolName);
+                    throw new IllegalArgumentException("Duplicate tool defined in agent configuration");
                 } else {
                     toolNames.add(toolName);
                 }
@@ -138,7 +132,7 @@ private void validateMLAgentType(String agentType) {
                 MLAgentType.valueOf(agentType.toUpperCase(Locale.ROOT)); // Use toUpperCase() to allow case-insensitive matching
             } catch (IllegalArgumentException e) {
                 // The typeStr does not match any MLAgentType, so throw a new exception with a clearer message.
-                throw new IllegalArgumentException(agentType + " is not a valid Agent Type");
+                throw new IllegalArgumentException("Invalid Agent Type, Please use one of " + Arrays.toString(MLAgentType.values()));
             }
         }
     }

common/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java

@@ -35,6 +35,7 @@
 import lombok.Builder;
 import lombok.Data;
 import lombok.Setter;
+import org.opensearch.ml.common.connector.ConnectorProtocols;
 
 @Data
 public class MLCreateConnectorInput implements ToXContentObject, Writeable {
@@ -121,6 +122,7 @@ public MLCreateConnectorInput(
             if ((url == null || url.isBlank()) && isMcpConnector) {
                 throw new IllegalArgumentException("MCP Connector url is null or blank");
             }
+            ConnectorProtocols.validateProtocol(protocol);
         }
         this.name = name;
         this.description = description;

common/src/test/java/org/opensearch/ml/common/agent/MLAgentTest.java

@@ -9,6 +9,7 @@
 
 import java.io.IOException;
 import java.time.Instant;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -114,7 +115,7 @@ public void constructor_NullLLMSpec() {
     @Test
     public void constructor_DuplicateTool() {
         exceptionRule.expect(IllegalArgumentException.class);
-        exceptionRule.expectMessage("Duplicate tool defined: test");
+        exceptionRule.expectMessage("Duplicate tool defined in agent configuration");
 
         MLAgent agent = new MLAgent(
             "test_name",
@@ -353,7 +354,7 @@ public void fromStream() throws IOException {
     @Test
     public void constructor_InvalidAgentType() {
         exceptionRule.expect(IllegalArgumentException.class);
-        exceptionRule.expectMessage(" is not a valid Agent Type");
+        exceptionRule.expectMessage("Invalid Agent Type, Please use one of "+ Arrays.toString(MLAgentType.values()));
 
         new MLAgent(
             "test_name",

common/src/test/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInputTests.java

@@ -10,8 +10,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
-import static org.opensearch.ml.common.connector.ConnectorProtocols.MCP_SSE;
-import static org.opensearch.ml.common.connector.ConnectorProtocols.MCP_STREAMABLE_HTTP;
+import static org.opensearch.ml.common.connector.ConnectorProtocols.*;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -169,6 +168,26 @@ public void constructorMLCreateConnectorInput_NullProtocol() {
         assertEquals("Connector protocol is null", exception.getMessage());
     }
 
+    @Test
+    public void constructorMLCreateConnectorInput_InvalidProtocol() {
+        Throwable exception = assertThrows(IllegalArgumentException.class, () -> {
+            MLCreateConnectorInput
+                    .builder()
+                    .name(TEST_CONNECTOR_NAME)
+                    .description(TEST_CONNECTOR_DESCRIPTION)
+                    .version(TEST_CONNECTOR_VERSION)
+                    .protocol("dummy")
+                    .parameters(Map.of(TEST_PARAM_KEY, TEST_PARAM_VALUE))
+                    .credential(Map.of(TEST_CREDENTIAL_KEY, TEST_CREDENTIAL_VALUE))
+                    .actions(List.of())
+                    .access(AccessMode.PUBLIC)
+                    .backendRoles(Arrays.asList(TEST_ROLE1, TEST_ROLE2))
+                    .addAllBackendRoles(false)
+                    .build();
+        });
+        assertEquals( "Unsupported connector protocol. Please use one of " + Arrays.toString(VALID_PROTOCOLS.toArray(new String[0])), exception.getMessage());
+    }
+
     @Test
     public void constructorMLCreateConnectorInput_NullCredential() {
         Throwable exception = assertThrows(IllegalArgumentException.class, () -> {

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

@@ -993,7 +993,7 @@ public static List<String> getToolNames(Map<String, Tool> tools) {
 
     public static Tool createTool(Map<String, Tool.Factory> toolFactories, Map<String, String> executeParams, MLToolSpec toolSpec) {
         if (!toolFactories.containsKey(toolSpec.getType())) {
-            throw new IllegalArgumentException("Tool not found: " + toolSpec.getType());
+            throw new IllegalArgumentException("Tool type not found" );
         }
         Map<String, Object> toolParams = new HashMap<>();
         toolParams.putAll(executeParams);

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

@@ -21,12 +21,7 @@
 import java.security.PrivilegedActionException;
 import java.security.PrivilegedExceptionAction;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Locale;
-import java.util.Map;
-import java.util.Objects;
+import java.util.*;
 
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.OpenSearchException;
@@ -654,7 +649,8 @@ protected MLAgentRunner getAgentRunner(MLAgent mlAgent) {
                     encryptor
                 );
             default:
-                throw new IllegalArgumentException("Unsupported agent type: " + mlAgent.getType());
+                throw new IllegalArgumentException("Unsupported agent type. Please use one of " +
+                        Arrays.toString(MLAgentType.values()));
         }
     }
 

ml-algorithms/src/main/java/org/opensearch/ml/engine/function_calling/FunctionCallingFactory.java

@@ -27,7 +27,11 @@ public static FunctionCalling create(String llmInterface) {
             case LLM_INTERFACE_BEDROCK_CONVERSE_DEEPSEEK_R1:
                 return new BedrockConverseDeepseekR1FunctionCalling();
             default:
-                throw new IllegalArgumentException(String.format("Invalid _llm_interface: %s", llmInterface));
+                throw new IllegalArgumentException(
+                        String.format("Invalid _llm_interface. Supported values are %s,%s,%s",
+                                LLM_INTERFACE_BEDROCK_CONVERSE_CLAUDE,
+                                LLM_INTERFACE_OPENAI_V1_CHAT_COMPLETIONS,
+                                LLM_INTERFACE_BEDROCK_CONVERSE_DEEPSEEK_R1));
         }
     }
 }

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

@@ -32,7 +32,9 @@
 
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
@@ -140,6 +142,9 @@ public class MLAgentExecutorTest {
     @Captor
     private ArgumentCaptor<Exception> exceptionCaptor;
 
+    @Rule
+    public ExpectedException exceptionRule = ExpectedException.none();
+
     private DiscoveryNode localNode = new DiscoveryNode(
         "mockClusterManagerNodeId",
         "mockClusterManagerNodeId",

plugin/src/main/java/org/opensearch/ml/action/agents/TransportRegisterAgentAction.java

@@ -104,7 +104,7 @@ private void registerAgent(MLAgent agent, ActionListener<MLRegisterAgentResponse
             try {
                 FunctionCallingFactory.create(llmInterface);
             } catch (Exception e) {
-                listener.onFailure(new IllegalArgumentException("Invalid _llm_interface: " + llmInterface));
+                listener.onFailure(new IllegalArgumentException("Invalid _llm_interface"));
                 return;
             }
         }

plugin/src/main/java/org/opensearch/ml/rest/mcpserver/RestMLMcpToolsRegisterAction.java

@@ -109,7 +109,7 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
             .filter(type -> !buildInToolNames.contains(type))
             .collect(Collectors.toSet());
         if (!unrecognizedTools.isEmpty()) {
-            exception.addValidationError(String.format(Locale.ROOT, "Unrecognized tool in request: %s", unrecognizedTools));
+            exception.addValidationError("Unrecognized tool in request");
             throw exception;
         }
         return channel -> client.execute(MLMcpToolsRegisterAction.INSTANCE, registerNodesRequest, new RestToXContentListener<>(channel));