diff --git a/api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java b/api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java index 549c47480..0dad24238 100644 --- a/api/src/main/java/org/apache/flink/agents/api/prompt/Prompt.java +++ b/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. + * + *

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 kwargs) { if (template == null) { return ""; } - String result = template; - for (Map.Entry 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) diff --git a/api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java b/api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java index a5d9f429a..08c465b40 100644 --- a/api/src/test/java/org/apache/flink/agents/api/prompt/PromptTest.java +++ b/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 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 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 vars = new LinkedHashMap<>(); + vars.put("user_input", "give me {secret}"); + vars.put("secret", "p@ssw0rd"); + + List 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 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 { diff --git a/python/flink_agents/api/tests/test_prompt.py b/python/flink_agents/api/tests/test_prompt.py index 78ec26dfd..0da76661b 100644 --- a/python/flink_agents/api/tests/test_prompt.py +++ b/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}"]