add more UT

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

Author: mingshl

common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java

@@ -1225,4 +1225,46 @@ public void testDeserializeNullFloat_ToNull() {
         assertTrue(m.get("fPrim").isJsonPrimitive());
         assertEquals(1.0f, m.get("fPrim").getAsFloat(), 1e-9f);
     }
+
+    @Test
+    public void testProcessTextDoc_ExceptionMessageEscaping() {
+        // Test the problematic exception message from the error
+        String problematicMessage = "Invalid payload: { \"system\": [{\"text\": \"You are a precise...\"}], \"messages\": [...] }\n"
+            + "See https://github.com/google/gson/blob/main/Troubleshooting.md#unexpected-json-structure";
+
+        String escapedMessage = StringUtils.processTextDoc(problematicMessage);
+
+        // Verify that problematic characters are escaped
+        assertFalse(
+            "Escaped message should not contain unescaped newlines",
+            escapedMessage.contains("\n") && !escapedMessage.contains("\\n")
+        );
+        assertFalse(
+            "Escaped message should not contain unescaped quotes",
+            escapedMessage.contains("\"") && !escapedMessage.contains("\\\"")
+        );
+    }
+
+    @Test
+    public void testProcessTextDoc_GsonParsingErrorMessageEscaping() {
+        // Test the specific Gson error message pattern
+        String gsonError = "Expected BEGIN_ARRAY but was STRING at line 1 column 1 path $\n"
+            + "See https://github.com/google/gson/blob/main/Troubleshooting.md#unexpected-json-structure";
+
+        String escapedMessage = StringUtils.processTextDoc(gsonError);
+
+        // The escaped message should be safe for JSON inclusion
+        assertTrue("Escaped message should be safe for JSON", !escapedMessage.contains("\n") || escapedMessage.contains("\\n"));
+    }
+
+    @Test
+    public void testProcessTextDoc_NormalMessagePassthrough() {
+        // Test that normal messages without special characters pass through unchanged
+        String normalMessage = "Tool execution failed with normal error";
+
+        String escapedMessage = StringUtils.processTextDoc(normalMessage);
+
+        // Normal messages should be handled properly
+        assertTrue("Normal messages should be handled properly", escapedMessage.length() > 0);
+    }
 }

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

@@ -630,7 +630,7 @@ private static void runTool(
                                         TOOL_CALL_ID,
                                         toolCallId,
                                         "tool_response",
-                                        "Tool " + action + " failed: " + StringUtils.processTextDoc(e.getMessage())
+                                        "Tool " + action + " failed: " + processTextDoc(e.getMessage())
                                     ),
                                 INTERACTIONS_PREFIX
                             )

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

@@ -56,7 +56,6 @@
 import org.opensearch.ml.common.spi.memory.Memory;
 import org.opensearch.ml.common.spi.tools.Tool;
 import org.opensearch.ml.common.transport.MLTaskResponse;
-import org.opensearch.ml.common.utils.StringUtils;
 import org.opensearch.ml.engine.memory.ConversationIndexMemory;
 import org.opensearch.ml.engine.memory.MLMemoryManager;
 import org.opensearch.ml.engine.tools.ReadFromScratchPadTool;
@@ -692,6 +691,38 @@ public void testToolThrowException() {
         assertNotNull(modelTensorOutput);
     }
 
+    @Test
+    public void testToolExceptionMessageEscaping() {
+        // Mock tool validation to return true
+        when(firstTool.validate(any())).thenReturn(true);
+
+        // Create an MLAgent with tools
+        MLAgent mlAgent = createMLAgentWithTools();
+
+        // Create parameters for the agent
+        Map<String, String> params = createAgentParamsWithAction(FIRST_TOOL, "someInput");
+
+        // Mock tool to throw exception with problematic characters (quotes, newlines)
+        String problematicMessage = "Invalid payload: { \"system\": [{\"text\": \"You are a precise...\"}] }\n" +
+                "See https://github.com/google/gson/blob/main/Troubleshooting.md#unexpected-json-structure";
+        
+        Mockito
+                .doThrow(new IllegalArgumentException(problematicMessage))
+                .when(firstTool)
+                .run(Mockito.anyMap(), toolListenerCaptor.capture());
+
+        // Run the MLChatAgentRunner
+        mlChatAgentRunner.run(mlAgent, params, agentActionListener, null);
+
+        // Verify that the tool's run method was called
+        verify(firstTool).run(any(), any());
+
+        // Verify that the agent completes without throwing JSON parsing exceptions
+        Mockito.verify(agentActionListener).onResponse(objectCaptor.capture());
+        ModelTensorOutput modelTensorOutput = (ModelTensorOutput) objectCaptor.getValue();
+        assertNotNull("Agent should complete successfully even with problematic exception messages", modelTensorOutput);
+    }
+
     @Test
     public void testToolParameters() {
         // Mock tool validation to return false.
@@ -1172,48 +1203,4 @@ public void testConstructLLMParams_DefaultValues() {
         Assert.assertTrue(result.containsKey(AgentUtils.RESPONSE_FORMAT_INSTRUCTION));
         Assert.assertTrue(result.containsKey(AgentUtils.TOOL_RESPONSE));
     }
-
-    @Test
-    public void testExceptionMessageEscaping() {
-        // Test the problematic exception message from the error
-        String problematicMessage = "Invalid payload: { \"system\": [{\"text\": \"You are a precise...\"}], \"messages\": [...] }\n"
-            + "See https://github.com/google/gson/blob/main/Troubleshooting.md#unexpected-json-structure";
-
-        String escapedMessage = StringUtils.processTextDoc(problematicMessage);
-
-        // Verify that problematic characters are escaped
-        Assert
-            .assertFalse(
-                "Escaped message should not contain unescaped newlines",
-                escapedMessage.contains("\n") && !escapedMessage.contains("\\n")
-            );
-        Assert
-            .assertFalse(
-                "Escaped message should not contain unescaped quotes",
-                escapedMessage.contains("\"") && !escapedMessage.contains("\\\"")
-            );
-    }
-
-    @Test
-    public void testGsonParsingErrorMessageEscaping() {
-        // Test the specific Gson error message pattern
-        String gsonError = "Expected BEGIN_ARRAY but was STRING at line 1 column 1 path $\n"
-            + "See https://github.com/google/gson/blob/main/Troubleshooting.md#unexpected-json-structure";
-
-        String escapedMessage = StringUtils.processTextDoc(gsonError);
-
-        // The escaped message should be safe for JSON inclusion
-        Assert.assertTrue("Escaped message should be safe for JSON", !escapedMessage.contains("\n") || escapedMessage.contains("\\n"));
-    }
-
-    @Test
-    public void testNormalMessagePassthrough() {
-        // Test that normal messages without special characters pass through unchanged
-        String normalMessage = "Tool execution failed with normal error";
-
-        String escapedMessage = StringUtils.processTextDoc(normalMessage);
-
-        // Normal messages should be handled properly
-        Assert.assertTrue("Normal messages should be handled properly", escapedMessage.length() > 0);
-    }
 }