From d1404a1299043fde789586f36770cd87f91ea149 Mon Sep 17 00:00:00 2001 From: Mingshi Liu Date: Sun, 9 Nov 2025 19:23:40 -0800 Subject: [PATCH 1/2] fix error message not proper escape Signed-off-by: Mingshi Liu --- .../algorithms/agent/MLChatAgentRunner.java | 8 +++- .../agent/MLChatAgentRunnerTest.java | 45 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java index 103f3f89b3..c12146c13b 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java @@ -625,7 +625,13 @@ private static void runTool( .add( substitute( tmpParameters.get(INTERACTION_TEMPLATE_TOOL_RESPONSE), - Map.of(TOOL_CALL_ID, toolCallId, "tool_response", "Tool " + action + " failed: " + e.getMessage()), + Map + .of( + TOOL_CALL_ID, + toolCallId, + "tool_response", + "Tool " + action + " failed: " + StringUtils.processTextDoc(e.getMessage()) + ), INTERACTIONS_PREFIX ) ); diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java index f6c3e3618e..6dbb87b684 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java @@ -56,6 +56,7 @@ 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; @@ -1171,4 +1172,48 @@ 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); + } } From 168c496a76a591254aeb14c55f6a954813304e23 Mon Sep 17 00:00:00 2001 From: Mingshi Liu Date: Mon, 17 Nov 2025 14:19:06 -0800 Subject: [PATCH 2/2] add more UT Signed-off-by: Mingshi Liu --- .../ml/common/utils/StringUtilsTest.java | 42 ++++++++++ .../algorithms/agent/MLChatAgentRunner.java | 2 +- .../agent/MLChatAgentRunnerTest.java | 77 ++++++++----------- 3 files changed, 75 insertions(+), 46 deletions(-) diff --git a/common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java b/common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java index e81ccc54a3..a4e65ce381 100644 --- a/common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java +++ b/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); + } } diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java index c12146c13b..09fcf67dda 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java +++ b/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 ) diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java index 6dbb87b684..acdd5d1795 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java +++ b/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 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); - } }