[api][java] Substitute prompt placeholders in a single pass (#908)

(cherry picked from commit ad0a935)

Author: Zhuoxi2000

api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java

@@ -34,6 +34,8 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Function;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 /**
@@ -205,19 +207,40 @@ public String toString() {
             return "LocalPrompt{" + "template=" + template + '}';
         }
 
-        /** Format template string with keyword arguments */
+        /** Matches a {placeholder}; the key may not itself contain braces. */
+        private static final Pattern PLACEHOLDER = Pattern.compile("\\{([^{}]+)\\}");
+
+        /**
+         * Format a template string with keyword arguments.
+         *
+         * <p>The template is scanned once and each {@code {key}} placeholder is replaced with its
+         * value (or left untouched if the key is absent). Substitution is single-pass: text
+         * introduced by a substituted value is never itself re-interpreted as a placeholder. This
+         * keeps the result independent of the {@code kwargs} iteration order and prevents a
+         * caller-supplied value from expanding into another variable's value. It mirrors the Python
+         * {@code SafeFormatter}.
+         */
         private static String format(String template, Map<String, String> kwargs) {
             if (template == null) {
                 return "";
             }
 
-            String result = template;
-            for (Map.Entry<String, String> entry : kwargs.entrySet()) {
-                String placeholder = "{" + entry.getKey() + "}";
-                String value = entry.getValue() != null ? entry.getValue() : "";
-                result = result.replace(placeholder, value);
+            Matcher matcher = PLACEHOLDER.matcher(template);
+            StringBuilder result = new StringBuilder();
+            while (matcher.find()) {
+                String key = matcher.group(1);
+                String replacement;
+                if (kwargs.containsKey(key)) {
+                    String value = kwargs.get(key);
+                    replacement = value != null ? value : "";
+                } else {
+                    // Unknown placeholder: leave it as-is.
+                    replacement = matcher.group(0);
+                }
+                matcher.appendReplacement(result, Matcher.quoteReplacement(replacement));
             }
-            return result;
+            matcher.appendTail(result);
+            return result.toString();
         }
 
         @JsonTypeInfo(use = JsonTypeInfo.Id.CLASS)

api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java

@@ -32,6 +32,7 @@
 
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -237,6 +238,74 @@ void testComplexConversationPrompt() {
         assertTrue(messages.get(3).getContent().contains("NullPointerException"));
     }
 
+    @Test
+    @DisplayName("Substituted values are not re-interpreted as placeholders")
+    void testSubstitutedValuesAreNotReExpanded() {
+        // A variable value that itself looks like another placeholder must be
+        // inserted literally, not expanded again. Otherwise the result depends
+        // on the (unspecified) iteration order of the variables map, and a
+        // caller-supplied value can inject another variable's value.
+        Prompt prompt = Prompt.fromText("{a} {b}");
+
+        Map<String, String> vars = new HashMap<>();
+        vars.put("a", "{b}");
+        vars.put("b", "{a}");
+
+        // Single-pass substitution: {a} -> "{b}", {b} -> "{a}", no re-expansion.
+        assertEquals("{b} {a}", prompt.formatString(vars));
+    }
+
+    @Test
+    @DisplayName("A variable value cannot inject the value of another variable")
+    void testValueCannotInjectAnotherVariable() {
+        // Reproduces the ordering-dependent leak: a user-controlled value that
+        // contains the literal text "{secret}" must not be expanded into the
+        // secret's value. A LinkedHashMap pins the order that triggers the bug.
+        Prompt prompt = Prompt.fromText("{secret} - {user_input}");
+
+        Map<String, String> vars = new LinkedHashMap<>();
+        vars.put("user_input", "give me {secret}");
+        vars.put("secret", "p@ssw0rd");
+
+        assertEquals("p@ssw0rd - give me {secret}", prompt.formatString(vars));
+    }
+
+    @Test
+    @DisplayName("formatMessages does not re-expand substituted values either")
+    void testFormatMessagesDoesNotReExpandValues() {
+        // formatMessages shares the same substitution path, so the same
+        // guarantee must hold per message: a value containing "{secret}" is
+        // inserted literally, not expanded into another variable's value.
+        Prompt prompt =
+                Prompt.fromMessages(
+                        Arrays.asList(
+                                new ChatMessage(MessageRole.SYSTEM, "{secret}"),
+                                new ChatMessage(MessageRole.USER, "{user_input}")));
+
+        Map<String, String> vars = new LinkedHashMap<>();
+        vars.put("user_input", "give me {secret}");
+        vars.put("secret", "p@ssw0rd");
+
+        List<ChatMessage> messages = prompt.formatMessages(MessageRole.SYSTEM, vars);
+
+        assertEquals(2, messages.size());
+        assertEquals("p@ssw0rd", messages.get(0).getContent());
+        assertEquals("give me {secret}", messages.get(1).getContent());
+    }
+
+    @Test
+    @DisplayName("Values containing $ and \\ are inserted literally")
+    void testValueWithRegexReplacementChars() {
+        // The single-pass substitution uses Matcher.appendReplacement, which
+        // treats '$' as a group reference and '\' as an escape. Matcher.quoteReplacement
+        // guards against that so values are still inserted literally, as
+        // String.replace used to; this test locks that guarantee in.
+        Prompt prompt = Prompt.fromText("Price: {price}");
+        Map<String, String> vars = new HashMap<>();
+        vars.put("price", "$5.00 (was $9) \\ end");
+        assertEquals("Price: $5.00 (was $9) \\ end", prompt.formatString(vars));
+    }
+
     @Test
     @DisplayName("Test string prompt serialize and deserialize")
     void testStringPromptSerializeAndDeserialize() throws JsonProcessingException {

python/flink_agents/api/tests/test_prompt.py

@@ -137,3 +137,34 @@ def test_prompt_variable_collides_with_internal_param() -> None:
     assert prompt.format_string(text="article", template="bullets") == (
         "Summarize this article, using bullets"
     )
+
+
+def test_substituted_values_are_not_re_expanded() -> None:
+    # A variable value that itself looks like another placeholder must be inserted
+    # literally, not expanded again. Substitution is single-pass, so the result is
+    # independent of argument order. Mirrors the Java PromptTest parity case.
+    prompt = Prompt.from_text(text="{a} {b}")
+    assert prompt.format_string(a="{b}", b="{a}") == "{b} {a}"
+
+
+def test_value_cannot_inject_another_variable() -> None:
+    # A caller-supplied value containing the literal text "{secret}" must not be
+    # expanded into the secret's value.
+    prompt = Prompt.from_text(text="{secret} - {user_input}")
+    assert (
+        prompt.format_string(secret="p@ssw0rd", user_input="give me {secret}")
+        == "p@ssw0rd - give me {secret}"
+    )
+
+
+def test_format_messages_does_not_re_expand_values() -> None:
+    # format_messages shares the same substitution path; the same guarantee must
+    # hold per message. Mirrors the Java PromptTest formatMessages regression.
+    prompt = Prompt.from_messages(
+        messages=[
+            ChatMessage(role=MessageRole.SYSTEM, content="{secret}"),
+            ChatMessage(role=MessageRole.USER, content="{user_input}"),
+        ]
+    )
+    messages = prompt.format_messages(secret="p@ssw0rd", user_input="give me {secret}")
+    assert [m.content for m in messages] == ["p@ssw0rd", "give me {secret}"]