From 14bcb074adad9deb9e658fdf2824625395e1bd87 Mon Sep 17 00:00:00 2001
From: Hailong Cui <ihailong@amazon.com>
Date: Thu, 18 Sep 2025 10:43:44 +0800
Subject: [PATCH 1/3] add escape output parameter

Signed-off-by: Hailong Cui <ihailong@amazon.com>
---
 .../org/opensearch/ml/common/utils/StringUtils.java   |  9 +++++++--
 .../org/opensearch/ml/common/utils/ToolUtils.java     |  1 +
 .../opensearch/ml/common/utils/StringUtilsTest.java   | 11 +++++++++--
 .../agent/MLConversationalFlowAgentRunner.java        |  8 +++++++-
 .../ml/engine/algorithms/agent/MLFlowAgentRunner.java | 11 ++++++++++-
 5 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java b/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java
index 61a275e5f6..356176d6c7 100644
--- a/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java
+++ b/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java
@@ -129,11 +129,16 @@ public static boolean isJson(String json) {
      *   prepareJsonValue("{\"key\":123}") → {\"key\":123} (valid JSON object, unchanged)
      * </pre>
      * @param input
+     * @param escape
      * @return
      */
-    public static String prepareJsonValue(String input) {
+    public static String prepareJsonValue(String input, boolean escape) {
         if (isJson(input)) {
-            return input;
+            if (!escape) {
+                return input;
+            } else {
+                return escapeJson(input);
+            }
         }
         return escapeJson(input);
     }
diff --git a/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java b/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java
index becf53e3c8..d7ee79cbd1 100644
--- a/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java
+++ b/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java
@@ -33,6 +33,7 @@
 public class ToolUtils {
 
     public static final String TOOL_OUTPUT_FILTERS_FIELD = "output_filter";
+    public static final String TOOL_OUTPUT_ESCAPED = "output_escaped";
     public static final String TOOL_REQUIRED_PARAMS = "required_parameters";
     public static final String NO_ESCAPE_PARAMS = "no_escape_params";
 
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 155f78d3d7..2013f4581e 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
@@ -857,14 +857,21 @@ public void testValidateFields_InvalidCharacterSet() {
     @Test
     public void prepareJsonValue_returnsRawIfJson() {
         String json = "{\"key\": 123}";
-        String result = StringUtils.prepareJsonValue(json);
+        String result = StringUtils.prepareJsonValue(json, false);
         assertSame(json, result);  // branch where isJson(input)==true
     }
 
+    @Test
+    public void prepareJsonValue_returnEscapeJsonIfForce() {
+        String json = "{\"key\": 123}";
+        String result = StringUtils.prepareJsonValue(json, true);
+        assertEquals("{\\\"key\\\": 123}", result);
+    }
+
     @Test
     public void prepareJsonValue_escapesBadCharsOtherwise() {
         String input = "Tom & Jerry \"<script>";
-        String escaped = StringUtils.prepareJsonValue(input);
+        String escaped = StringUtils.prepareJsonValue(input, false);
         assertNotEquals(input, escaped);
         assertFalse(StringUtils.isJson(escaped));
         assertEquals("Tom & Jerry \\\"<script>", escaped);
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java
index 5a4ce8e796..19a6fbb184 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java
@@ -9,6 +9,7 @@
 import static org.opensearch.ml.common.conversation.ActionConstants.AI_RESPONSE_FIELD;
 import static org.opensearch.ml.common.conversation.ActionConstants.MEMORY_ID;
 import static org.opensearch.ml.common.conversation.ActionConstants.PARENT_INTERACTION_ID_FIELD;
+import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_ESCAPED;
 import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_FILTERS_FIELD;
 import static org.opensearch.ml.common.utils.ToolUtils.filterToolOutput;
 import static org.opensearch.ml.common.utils.ToolUtils.getToolName;
@@ -262,7 +263,12 @@ private void processOutput(
         String outputKey = toolName + ".output";
         Map<String, String> toolParameters = ToolUtils.buildToolParameters(params, previousToolSpec, tenantId);
         String filteredOutput = parseResponse(filterToolOutput(toolParameters, output));
-        params.put(outputKey, StringUtils.prepareJsonValue(filteredOutput));
+        params
+            .put(
+                outputKey,
+                StringUtils
+                    .prepareJsonValue(filteredOutput, Boolean.parseBoolean(toolParameters.getOrDefault(TOOL_OUTPUT_ESCAPED, "false")))
+            );
         boolean traceDisabled = params.containsKey(DISABLE_TRACE) && Boolean.parseBoolean(params.get(DISABLE_TRACE));
 
         if (previousToolSpec.isIncludeOutputInAgentResponse() || finalI == toolSpecs.size()) {
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
index 4708b4a8e4..7be9603520 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.ml.engine.algorithms.agent;
 
+import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_ESCAPED;
 import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_FILTERS_FIELD;
 import static org.opensearch.ml.common.utils.ToolUtils.filterToolOutput;
 import static org.opensearch.ml.common.utils.ToolUtils.getToolName;
@@ -114,7 +115,15 @@ public void run(MLAgent mlAgent, Map<String, String> params, ActionListener<Obje
                     String outputKey = toolName + ".output";
                     Map<String, String> toolParameters = ToolUtils.buildToolParameters(params, previousToolSpec, mlAgent.getTenantId());
                     String filteredOutput = parseResponse(filterToolOutput(toolParameters, output));
-                    params.put(outputKey, StringUtils.prepareJsonValue(filteredOutput));
+                    params
+                        .put(
+                            outputKey,
+                            StringUtils
+                                .prepareJsonValue(
+                                    filteredOutput,
+                                    Boolean.parseBoolean(toolParameters.getOrDefault(TOOL_OUTPUT_ESCAPED, "false"))
+                                )
+                        );
                     if (previousToolSpec.isIncludeOutputInAgentResponse() || finalI == toolSpecs.size()) {
                         if (toolParameters.containsKey(TOOL_OUTPUT_FILTERS_FIELD)) {
                             flowAgentOutput.add(ModelTensor.builder().name(outputKey).result(filteredOutput).build());

From 97c269ad15c335623ee5626b1ec58eb009c2efc1 Mon Sep 17 00:00:00 2001
From: Hailong Cui <ihailong@amazon.com>
Date: Wed, 24 Sep 2025 10:08:03 +0800
Subject: [PATCH 2/3] rename parameter name

Signed-off-by: Hailong Cui <ihailong@amazon.com>
---
 .../main/java/org/opensearch/ml/common/utils/ToolUtils.java  | 2 +-
 .../algorithms/agent/MLConversationalFlowAgentRunner.java    | 5 ++---
 .../ml/engine/algorithms/agent/MLFlowAgentRunner.java        | 4 ++--
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java b/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java
index d7ee79cbd1..d80b8c6794 100644
--- a/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java
+++ b/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java
@@ -33,7 +33,7 @@
 public class ToolUtils {
 
     public static final String TOOL_OUTPUT_FILTERS_FIELD = "output_filter";
-    public static final String TOOL_OUTPUT_ESCAPED = "output_escaped";
+    public static final String TOOL_ESCAPE_OUTPUT = "escape_output";
     public static final String TOOL_REQUIRED_PARAMS = "required_parameters";
     public static final String NO_ESCAPE_PARAMS = "no_escape_params";
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java
index 19a6fbb184..1983acef0e 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLConversationalFlowAgentRunner.java
@@ -9,7 +9,7 @@
 import static org.opensearch.ml.common.conversation.ActionConstants.AI_RESPONSE_FIELD;
 import static org.opensearch.ml.common.conversation.ActionConstants.MEMORY_ID;
 import static org.opensearch.ml.common.conversation.ActionConstants.PARENT_INTERACTION_ID_FIELD;
-import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_ESCAPED;
+import static org.opensearch.ml.common.utils.ToolUtils.TOOL_ESCAPE_OUTPUT;
 import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_FILTERS_FIELD;
 import static org.opensearch.ml.common.utils.ToolUtils.filterToolOutput;
 import static org.opensearch.ml.common.utils.ToolUtils.getToolName;
@@ -266,8 +266,7 @@ private void processOutput(
         params
             .put(
                 outputKey,
-                StringUtils
-                    .prepareJsonValue(filteredOutput, Boolean.parseBoolean(toolParameters.getOrDefault(TOOL_OUTPUT_ESCAPED, "false")))
+                StringUtils.prepareJsonValue(filteredOutput, Boolean.parseBoolean(toolParameters.getOrDefault(TOOL_ESCAPE_OUTPUT, "false")))
             );
         boolean traceDisabled = params.containsKey(DISABLE_TRACE) && Boolean.parseBoolean(params.get(DISABLE_TRACE));
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
index 7be9603520..e0da275c31 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
@@ -5,7 +5,7 @@
 
 package org.opensearch.ml.engine.algorithms.agent;
 
-import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_ESCAPED;
+import static org.opensearch.ml.common.utils.ToolUtils.TOOL_ESCAPE_OUTPUT;
 import static org.opensearch.ml.common.utils.ToolUtils.TOOL_OUTPUT_FILTERS_FIELD;
 import static org.opensearch.ml.common.utils.ToolUtils.filterToolOutput;
 import static org.opensearch.ml.common.utils.ToolUtils.getToolName;
@@ -121,7 +121,7 @@ public void run(MLAgent mlAgent, Map<String, String> params, ActionListener<Obje
                             StringUtils
                                 .prepareJsonValue(
                                     filteredOutput,
-                                    Boolean.parseBoolean(toolParameters.getOrDefault(TOOL_OUTPUT_ESCAPED, "false"))
+                                    Boolean.parseBoolean(toolParameters.getOrDefault(TOOL_ESCAPE_OUTPUT, "false"))
                                 )
                         );
                     if (previousToolSpec.isIncludeOutputInAgentResponse() || finalI == toolSpecs.size()) {

From 1192c34f2ac3959dbf8dccea055a114da7641486 Mon Sep 17 00:00:00 2001
From: Hailong Cui <ihailong@amazon.com>
Date: Thu, 25 Sep 2025 07:40:00 +0800
Subject: [PATCH 3/3] update parameter name

Signed-off-by: Hailong Cui <ihailong@amazon.com>
---
 .../java/org/opensearch/ml/common/utils/StringUtils.java    | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java b/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java
index 467ce7665c..b70b7a1027 100644
--- a/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java
+++ b/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java
@@ -139,12 +139,12 @@ public static boolean isJson(String json) {
      *   prepareJsonValue("{\"key\":123}") → {\"key\":123} (valid JSON object, unchanged)
      * </pre>
      * @param input
-     * @param escape
+     * @param escapeJson
      * @return
      */
-    public static String prepareJsonValue(String input, boolean escape) {
+    public static String prepareJsonValue(String input, boolean escapeJson) {
         if (isJson(input)) {
-            if (!escape) {
+            if (!escapeJson) {
                 return input;
             } else {
                 return escapeJson(input);