diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py index b640fe00b7e..7b346f12a4b 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py @@ -31,7 +31,7 @@ import re import sys import uuid -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from dataclasses import dataclass, field from decimal import Decimal as _Decimal from enum import Enum @@ -67,6 +67,73 @@ _SAFE_PATH_SEGMENT_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$") +def _skip_powerfx_opaque_token(formula: str, start: int) -> int: + """Return the end of an ordinary quoted token or comment, or start if neither.""" + quote = formula[start] + if quote in ('"', "'"): + pos = start + 1 + while pos < len(formula): + if formula[pos] == quote: + if pos + 1 < len(formula) and formula[pos + 1] == quote: + pos += 2 + continue + return pos + 1 + pos += 1 + return pos + if formula.startswith("//", start): + pos = start + 2 + while pos < len(formula) and formula[pos] not in "\r\n": + pos += 1 + return pos + if formula.startswith("/*", start): + end = formula.find("*/", start + 2) + return len(formula) if end == -1 else end + 2 + return start + + +def _iter_powerfx_expression_indices(formula: str) -> Iterator[int]: + """Yield code positions, excluding literals/comments but including interpolation expressions.""" + # None denotes interpolated text; integers track record braces in expression sections. + scopes: list[int | None] = [0] + cursor = 0 + while cursor < len(formula): + depth = scopes[-1] + char = formula[cursor] + if depth is None: + if formula.startswith(('""', "{{", "}}"), cursor): + cursor += 2 + continue + if char == "{": + scopes.append(0) + elif char == '"': + scopes.pop() + cursor += 1 + continue + + if formula.startswith('$"', cursor): + scopes.append(None) + cursor += 2 + continue + + token_end = _skip_powerfx_opaque_token(formula, cursor) + if token_end != cursor: + cursor = token_end + continue + + if char == "{": + scopes[-1] = depth + 1 + elif char == "}": + if depth > 0: + scopes[-1] = depth - 1 + elif len(scopes) > 1: + scopes.pop() + cursor += 1 + continue + + yield cursor + cursor += 1 + + @dataclass(frozen=True) class DeclarativeEnvConfig: """Configuration that populates the PowerFx ``Env`` symbol for a workflow. @@ -667,11 +734,11 @@ def _preprocess_custom_functions(self, formula: str, temp_writes: list[tuple[str When they appear nested inside other functions (e.g., Upper(MessageText(...))), we need to evaluate them first and replace with the result. - For long strings (>500 chars), the result is stored in a temporary state variable - to avoid exceeding PowerFx's 1000 character expression limit. This is a limitation - of the Python PowerFx wrapper (powerfx package), which doesn't expose the - MaximumExpressionLength configuration that the .NET PowerFxConfig provides. - The .NET implementation defaults to 10,000 characters, while Python defaults to 1,000. + Results are stored in temporary state variables so untrusted message text is + passed to PowerFx as data rather than inserted into formula source. + Temporary names avoid existing Local keys and references in the original formula. + Literal text, quoted identifiers, and comments are left untouched; expression + sections inside PowerFx interpolated strings are preprocessed as code. Args: formula: The PowerFx formula to pre-process @@ -684,90 +751,53 @@ def _preprocess_custom_functions(self, formula: str, temp_writes: list[tuple[str Returns: The rewritten formula. """ - import re - - # Threshold for storing in state vs embedding as literal. - # The Python PowerFx wrapper defaults to a 1000 char expression limit (vs 10,000 in .NET). - # We use 500 to leave room for the rest of the expression around the replaced value. - MAX_INLINE_LENGTH = 500 - temp_var_counter = 0 - - # Custom functions that need pre-processing: (regex pattern, handler) - custom_functions = [ - (r"MessageText\(", self._eval_and_replace_message_text), - ] - - for pattern, handler in custom_functions: - # Find all occurrences of the custom function - while True: - match = re.search(pattern, formula) - if not match: + reserved_names = {name.casefold() for name in self.get_state_data().get("Local", {})} + # Reserve formula references too, so previously undefined names stay undefined. + folded_formula = formula.casefold() + function_name = "MessageText" + call_prefix = f"{function_name}(" + result: list[str] = [] + copied_until = 0 + positions = _iter_powerfx_expression_indices(formula) + + for cursor in positions: + if not formula.startswith(call_prefix, cursor): + continue + + paren_start = cursor + len(function_name) + depth = 1 + for pos in positions: + if pos <= paren_start: + continue + char = formula[pos] + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: break + else: + break - # Find the matching closing parenthesis - start = match.start() - paren_start = match.end() - 1 # Position of opening ( - depth = 1 - pos = paren_start + 1 - in_string = False - escape_next = False - - while pos < len(formula) and depth > 0: - char = formula[pos] - if escape_next: - escape_next = False - pos += 1 - continue - if char == "\\": - escape_next = True - pos += 1 - continue - if char == '"' and not escape_next: - in_string = not in_string - elif not in_string: - if char == "(": - depth += 1 - elif char == ")": - depth -= 1 - pos += 1 - - if depth != 0: - # Malformed expression, skip + inner_expr = formula[paren_start + 1 : pos] + replacement = self._eval_and_replace_message_text(inner_expr) + while True: + temp_var_name = f"_TempMessageText{temp_var_counter}" + temp_var_counter += 1 + folded_name = temp_var_name.casefold() + if folded_name not in reserved_names and folded_name not in folded_formula: break - - # Extract the inner expression (between parentheses) - end = pos - inner_expr = formula[paren_start + 1 : end - 1] - - # Evaluate and get replacement - replacement = handler(inner_expr) - - # Replace in formula - if isinstance(replacement, str): - if len(replacement) > MAX_INLINE_LENGTH: - # Store long results in an underscore-prefixed temp key; - # record the prior value so eval() can restore it. - temp_var_name = f"_TempMessageText{temp_var_counter}" - temp_var_counter += 1 - temp_var_path = f"Local.{temp_var_name}" - temp_writes.append((temp_var_path, self.get(temp_var_path, default=self._MISSING))) - self.set(temp_var_path, replacement) - replacement_str = temp_var_path - logger.debug( - f"Stored long MessageText result ({len(replacement)} chars) " - f"in temp variable {temp_var_name}" - ) - else: - # Short strings can be embedded directly - escaped = replacement.replace('"', '""') - replacement_str = f'"{escaped}"' - else: - replacement_str = str(replacement) if replacement is not None else '""' - - formula = formula[:start] + replacement_str + formula[end:] - - return formula + temp_var_path = f"Local.{temp_var_name}" + temp_writes.append((temp_var_path, self.get(temp_var_path, default=self._MISSING))) + self.set(temp_var_path, replacement) + result.append(formula[copied_until:cursor]) + result.append(temp_var_path) + logger.debug(f"Stored MessageText result ({len(replacement)} chars) in temp variable {temp_var_name}") + copied_until = pos + 1 + + result.append(formula[copied_until:]) + return "".join(result) def _eval_and_replace_message_text(self, inner_expr: str) -> str: """Evaluate MessageText() and return the text result. diff --git a/python/packages/declarative/tests/test_declarative_state_path_safety.py b/python/packages/declarative/tests/test_declarative_state_path_safety.py index 2446fc3cf4d..40683e5e2ac 100644 --- a/python/packages/declarative/tests/test_declarative_state_path_safety.py +++ b/python/packages/declarative/tests/test_declarative_state_path_safety.py @@ -301,6 +301,165 @@ def test_end_to_end_send_activity_payload_neutralized( # --------------------------------------------------------------------------- +class TestMessageTextPreprocessing: + def test_replacement_text_is_not_rescanned_as_formula( + self, + state: DeclarativeWorkflowState, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls: list[str] = [] + + def replace_message_text(inner_expr: str) -> str: + calls.append(inner_expr) + if inner_expr == "Local.Messages": + return "MessageText(Local.Secret)" + pytest.fail(f"Replacement text was re-scanned: {inner_expr}") + + monkeypatch.setattr(state, "_eval_and_replace_message_text", replace_message_text) + temp_writes: list[tuple[str, Any]] = [] + + formula = state._preprocess_custom_functions("Upper(MessageText(Local.Messages))", temp_writes) + + assert formula == "Upper(Local._TempMessageText0)" + assert calls == ["Local.Messages"] + assert state.get("Local._TempMessageText0") == "MessageText(Local.Secret)" + + def test_interpolation_preprocesses_only_original_expression_sections( + self, state: DeclarativeWorkflowState, monkeypatch: pytest.MonkeyPatch + ) -> None: + calls: list[str] = [] + message_text = '$"MessageText(Local.Secret) {MessageText(Local.Secret)}"' + + def replace_message_text(inner_expr: str) -> str: + calls.append(inner_expr) + return message_text + + monkeypatch.setattr(state, "_eval_and_replace_message_text", replace_message_text) + temp_writes: list[tuple[str, Any]] = [] + formula = '$"Literal MessageText(Local.Messages) {{MessageText(Local.Messages)}} {MessageText(Local.Messages)}"' + + result = state._preprocess_custom_functions(formula, temp_writes) + + assert result == ( + '$"Literal MessageText(Local.Messages) {{MessageText(Local.Messages)}} {Local._TempMessageText0}"' + ) + assert calls == ["Local.Messages"] + assert state.get("Local._TempMessageText0") == message_text + assert temp_writes == [("Local._TempMessageText0", state._MISSING)] + + def test_message_text_inside_string_literal_is_not_preprocessed( + self, + state: DeclarativeWorkflowState, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls: list[str] = [] + + def replace_message_text(inner_expr: str) -> str: + calls.append(inner_expr) + return "hello" + + monkeypatch.setattr(state, "_eval_and_replace_message_text", replace_message_text) + + formula = state._preprocess_custom_functions( + 'Concatenate("MessageText(Local.Secret)", MessageText(Local.Messages))', + [], + ) + + assert formula == 'Concatenate("MessageText(Local.Secret)", Local._TempMessageText0)' + assert calls == ["Local.Messages"] + + @pytest.mark.parametrize( + "formula", + [ + 'Upper("Say ""MessageText(ignored)""")', + "Upper(Local.'MessageText(ignored)')", + "Upper(Local.'Archived'' MessageText(ignored)')", + "1 // MessageText(ignored)", + "1 // MessageText(ignored)\n + 2", + "1 // MessageText(ignored)\r + 2", + "1 // MessageText(ignored)\r\n + 2", + "/* MessageText(ignored) */ 1", + "1 /* unmatched ) ' \" MessageText(ignored) */", + ], + ) + def test_quoted_tokens_and_comments_are_not_preprocessed( + self, state: DeclarativeWorkflowState, monkeypatch: pytest.MonkeyPatch, formula: str + ) -> None: + def unexpected_call(inner_expr: str) -> str: + pytest.fail(f"Processed a call inside a quoted token or comment: {inner_expr}") + + monkeypatch.setattr(state, "_eval_and_replace_message_text", unexpected_call) + temp_writes: list[tuple[str, Any]] = [] + + assert state._preprocess_custom_functions(formula, temp_writes) == formula + assert temp_writes == [] + + @pytest.mark.parametrize( + "inner_expr", + [ + "Local.'Messages) archive'", + "Local.'Messages'' ) archive'", + 'If(")""(" = ")""(", Local.Messages, Local.Messages)', + "Local.Messages /* ) ' \" MessageText(ignored) */", + "Local.Messages // ) MessageText(ignored)\n", + "Local.Messages // ) MessageText(ignored)\r", + "Local.Messages // ) MessageText(ignored)\r\n", + ], + ) + def test_call_boundaries_ignore_quoted_tokens_and_comments( + self, state: DeclarativeWorkflowState, monkeypatch: pytest.MonkeyPatch, inner_expr: str + ) -> None: + calls: list[str] = [] + + def replace_message_text(expression: str) -> str: + calls.append(expression) + return "hello" + + monkeypatch.setattr(state, "_eval_and_replace_message_text", replace_message_text) + + formula = state._preprocess_custom_functions(f"Upper(MessageText({inner_expr}))", []) + + assert formula == "Upper(Local._TempMessageText0)" + assert calls == [inner_expr] + + @pytest.mark.parametrize( + ("existing_name", "existing_value"), + [("_TempMessageText0", None), ("_tempmessagetext0", "user-value")], + ) + def test_temporary_binding_avoids_existing_names( + self, + state: DeclarativeWorkflowState, + monkeypatch: pytest.MonkeyPatch, + existing_name: str, + existing_value: str | None, + ) -> None: + state.set(f"Local.{existing_name}", existing_value) + monkeypatch.setattr(state, "_eval_and_replace_message_text", lambda inner_expr: "hello") + temp_writes: list[tuple[str, Any]] = [] + + formula = state._preprocess_custom_functions("Upper(MessageText(Local.Messages))", temp_writes) + + assert formula == "Upper(Local._TempMessageText1)" + assert state.get_state_data()["Local"] == {existing_name: existing_value, "_TempMessageText1": "hello"} + assert temp_writes == [("Local._TempMessageText1", state._MISSING)] + + @pytest.mark.parametrize( + "reference", ["Local._TempMessageText0", "Local.'_TempMessageText0'", "Local._tempmessagetext0"] + ) + def test_temporary_binding_avoids_formula_references( + self, + state: DeclarativeWorkflowState, + monkeypatch: pytest.MonkeyPatch, + reference: str, + ) -> None: + monkeypatch.setattr(state, "_eval_and_replace_message_text", lambda inner_expr: "hello") + + formula = state._preprocess_custom_functions(f"Upper(MessageText(Local.Messages) & {reference})", []) + + assert formula == f"Upper(Local._TempMessageText1 & {reference})" + assert "_TempMessageText0" not in state.get_state_data()["Local"] + + @_requires_powerfx class TestPowerFxStillWorks: def test_simple_powerfx_expression_evaluates(self, state: DeclarativeWorkflowState) -> None: @@ -308,6 +467,86 @@ def test_simple_powerfx_expression_evaluates(self, state: DeclarativeWorkflowSta state.set("Local.y", 7) assert state.eval("=Local.x * Local.y") == 42 + @pytest.mark.parametrize("name", ["Messages) archive", "Messages' ) archive"]) + def test_message_text_argument_with_quoted_identifier(self, state: DeclarativeWorkflowState, name: str) -> None: + state.set(f"Local.{name}", [{"text": "hello", "contents": [{"type": "text", "text": "hello"}]}]) + quoted_name = name.replace("'", "''") + original_local = state.get_state_data()["Local"].copy() + + assert state.eval(f"=Upper(MessageText(Local.'{quoted_name}'))") == "HELLO" + assert state.get_state_data()["Local"] == original_local + + @pytest.mark.parametrize( + ("formula", "expected"), + [ + ('$"Message: {MessageText(Local.Messages)}"', "Message: hello"), + ( + '$"Literal MessageText(ignored) ""{{{MessageText(Local.Messages)}}}"""', + 'Literal MessageText(ignored) "{hello}"', + ), + ( + '$"{With({nested: {text: $"Inner {MessageText(Local.Messages)}"}}, nested.text)}"', + "Inner hello", + ), + ('$"{MessageText(Local.Messages)}:{Upper(MessageText(Local.Other))}"', "hello:OTHER"), + ('$"{/* } MessageText(Local.Secret) */ MessageText(Local.Messages)}"', "hello"), + ( + ( + 'Upper(MessageText(If($"Literal ) {MessageText(Local.Messages)}" = "Literal ) hello", ' + "Local.Messages, Local.Other)))" + ), + "HELLO", + ), + ], + ids=[ + "basic", + "literal-text-and-escapes", + "nested-records-and-interpolation", + "multiple-expressions", + "commented-interpolation-delimiter", + "interpolation-in-function-argument", + ], + ) + def test_expression_sections_evaluate_with_powerfx( + self, state: DeclarativeWorkflowState, formula: str, expected: str + ) -> None: + """Only expression sections of interpolated strings may evaluate MessageText calls.""" + state.set("Local.Messages", [{"text": "hello", "contents": [{"type": "text", "text": "hello"}]}]) + state.set("Local.Other", [{"text": "other", "contents": [{"type": "text", "text": "other"}]}]) + original_local = state.get_state_data()["Local"].copy() + + assert state.eval(f"={formula}") == expected + assert state.get_state_data()["Local"] == original_local + + @pytest.mark.parametrize("name", ["MessageText(ignored)", "Archived' MessageText(ignored)"]) + def test_quoted_identifier_with_function_like_text(self, state: DeclarativeWorkflowState, name: str) -> None: + state.set(f"Local.{name}", "hello") + quoted_name = name.replace("'", "''") + original_local = state.get_state_data()["Local"].copy() + + assert state.eval(f"=Upper(Local.'{quoted_name}')") == "HELLO" + assert state.get_state_data()["Local"] == original_local + + @pytest.mark.parametrize( + ("formula", "expected"), + [ + ("1 // MessageText(1 / 0)", 1), + ("1 /* MessageText(1 / 0) */", 1), + ("/* MessageText(1 / 0) */ Upper(MessageText(Local.Messages))", "HELLO"), + ("// MessageText(1 / 0)\n Upper(MessageText(Local.Messages))", "HELLO"), + ("Upper(MessageText(Local.Messages /* ) MessageText(1 / 0) */))", "HELLO"), + ("Upper(MessageText(Local.Messages // ) MessageText(1 / 0)\n))", "HELLO"), + ], + ) + def test_message_text_in_comments_is_not_evaluated( + self, state: DeclarativeWorkflowState, formula: str, expected: str | int + ) -> None: + state.set("Local.Messages", [{"text": "hello", "contents": [{"type": "text", "text": "hello"}]}]) + original_local = state.get_state_data()["Local"].copy() + + assert state.eval(f"={formula}") == expected + assert state.get_state_data()["Local"] == original_local + def test_internal_temp_message_text_still_works(self, state: DeclarativeWorkflowState) -> None: """Long MessageText() results round-trip and the temp key is removed after eval.""" long_text = "A" * 600 @@ -323,25 +562,59 @@ def test_internal_temp_message_text_still_works(self, state: DeclarativeWorkflow remaining = sorted(k for k in local if k.startswith("_TempMessageText")) assert not remaining, f"Temporary keys remain in Local: {remaining}" - def test_message_text_eval_preserves_user_temp_value(self, state: DeclarativeWorkflowState) -> None: - """User state at the temp key path survives a long MessageText eval.""" - long_text = "A" * 600 + @pytest.mark.parametrize("text", ["hello", "A" * 600], ids=["short", "long"]) + def test_message_text_eval_preserves_user_temp_value(self, state: DeclarativeWorkflowState, text: str) -> None: + """Existing symbols retain their values during evaluation, not just after cleanup.""" state.set("Local._TempMessageText0", "user-important-value") state.set( "Local.Messages", - [{"text": long_text, "contents": [{"type": "text", "text": long_text}]}], + [{"text": text, "contents": [{"type": "text", "text": text}]}], ) + original_local = state.get_state_data()["Local"].copy() - result = state.eval("=Upper(MessageText(Local.Messages))") - assert result == "A" * 600 - assert state.get("Local._TempMessageText0") == "user-important-value" + result = state.eval("=Upper(Local._TempMessageText0 & MessageText(Local.Messages))") + + assert result == f"USER-IMPORTANT-VALUE{text.upper()}" + assert state.get_state_data()["Local"] == original_local + + def test_message_text_eval_does_not_define_missing_symbol(self, state: DeclarativeWorkflowState) -> None: + state.set("Local.Messages", [{"text": "hello", "contents": [{"type": "text", "text": "hello"}]}]) + original_local = state.get_state_data()["Local"].copy() + + result = state.eval("=Upper(MessageText(Local.Messages) & Local._TempMessageText0)") + + assert result is None + assert state.get_state_data()["Local"] == original_local + @pytest.mark.parametrize( + "second_argument", + ["Local.Second", 'If(Upper(MessageText(Local.First)) = "HELLO", Local.Second, Local.First)'], + ) + def test_multiple_message_text_calls_use_distinct_bindings( + self, state: DeclarativeWorkflowState, second_argument: str + ) -> None: + state.set("Local._TempMessageText0", "prefix") + state.set("Local._TempMessageText1", "suffix") + state.set("Local.First", [{"text": "hello", "contents": [{"type": "text", "text": "hello"}]}]) + state.set("Local.Second", [{"text": "world", "contents": [{"type": "text", "text": "world"}]}]) + original_local = state.get_state_data()["Local"].copy() + + result = state.eval( + "=Upper(Local._TempMessageText0 & MessageText(Local.First) & " + f"MessageText({second_argument}) & Local._TempMessageText1)" + ) + + assert result == "PREFIXHELLOWORLDSUFFIX" + assert state.get_state_data()["Local"] == original_local + + @pytest.mark.parametrize("text", ["hello", "A" * 600], ids=["short", "long"]) def test_message_text_eval_cleans_up_on_powerfx_failure( self, state: DeclarativeWorkflowState, - monkeypatch, + monkeypatch: pytest.MonkeyPatch, + text: str, ) -> None: - """Temp key is removed even when PowerFx evaluation raises.""" + """New temp keys are removed and existing state survives evaluation errors.""" from agent_framework_declarative._workflows import _declarative_base as base class _FailingEngine: @@ -350,15 +623,16 @@ def eval(self, *args: Any, **kwargs: Any) -> Any: monkeypatch.setattr(base, "Engine", _FailingEngine) - long_text = "A" * 600 + state.set("Local._TempMessageText0", "user-important-value") + state.set("Local._TempMessageText1", None) state.set( "Local.Messages", - [{"text": long_text, "contents": [{"type": "text", "text": long_text}]}], + [{"text": text, "contents": [{"type": "text", "text": text}]}], ) + monkeypatch.setattr(state, "_eval_and_replace_message_text", lambda inner_expr: text) + original_local = state.get_state_data()["Local"].copy() with pytest.raises(RuntimeError, match="boom"): state.eval("=Upper(MessageText(Local.Messages))") - local = state.get_state_data().get("Local", {}) - remaining = sorted(k for k in local if k.startswith("_TempMessageText")) - assert not remaining, f"Temporary keys remain in Local after PowerFx failure: {remaining}" + assert state.get_state_data()["Local"] == original_local diff --git a/python/packages/declarative/tests/test_graph_coverage.py b/python/packages/declarative/tests/test_graph_coverage.py index 43662c82c2c..2d0b5cae2b1 100644 --- a/python/packages/declarative/tests/test_graph_coverage.py +++ b/python/packages/declarative/tests/test_graph_coverage.py @@ -2762,8 +2762,8 @@ async def test_multiplication_normal(self, mock_state): class TestLongMessageTextHandling: """Tests for handling long MessageText results that exceed PowerFx limits.""" - async def test_short_message_text_embedded_inline(self, mock_state): - """Test that short MessageText results are embedded inline.""" + async def test_short_message_text_round_trips_without_residual_temp_state(self, mock_state): + """Test that short MessageText results are removed from temporary state after evaluation.""" state = DeclarativeWorkflowState(mock_state) state.initialize() @@ -2771,11 +2771,11 @@ async def test_short_message_text_embedded_inline(self, mock_state): short_text = "Hello world" state.set("Local.Messages", [{"text": short_text, "contents": [{"type": "text", "text": short_text}]}]) - # Evaluate a formula with MessageText - should embed inline + # Evaluate a formula with MessageText. result = state.eval("=Upper(MessageText(Local.Messages))") assert result == "HELLO WORLD" - # No temp variable should be created for short strings + # Temporary state should be cleaned up after evaluation. temp_var = state.get("Local._TempMessageText0") assert temp_var is None