Python: Bound declarative custom-function matches to one whole call

python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py

@@ -46,6 +46,10 @@
 from agent_framework._workflows._state import State
 
 from ._errors import DeclarativeWorkflowError
+from ._powerfx_functions import (
+    _iter_powerfx_expression_indices,  # pyright: ignore[reportPrivateUsage]
+    whole_call_args,
+)
 from ._powerfx_limits import _PowerFxStateBudget, _validate_powerfx_state  # pyright: ignore[reportPrivateUsage]
 from ._state_path import _is_safe_path_segment  # pyright: ignore[reportPrivateUsage]
 
@@ -68,73 +72,6 @@
 _ENV_REFERENCE_RE = re.compile(r"\bEnv\.([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.
@@ -721,14 +658,14 @@ def _eval_custom_function(self, formula: str) -> Any | None:
 
         Returns None if the formula is not a custom function call.
         """
-        import re
+        formula = formula.strip()
 
         # Concat/Concatenate - string concatenation
         # In standard PowerFx, Concatenate is for strings, Concat is for tables.
         # Copilot Studio uses Concat for strings, so we support both.
-        match = re.match(r"(?:Concat|Concatenate)\((.+)\)$", formula.strip())
-        if match:
-            args_str = match.group(1)
+        args_str = whole_call_args(formula, "Concat", "Concatenate")
+        # Empty arguments fall through so PowerFx reports the arity error.
+        if args_str:
             # Parse comma-separated arguments (handling nested parentheses)
             args = self._parse_function_args(args_str)
             evaluated_args: list[str] = []
@@ -747,24 +684,24 @@ def _eval_custom_function(self, formula: str) -> Any | None:
             return "".join(evaluated_args)
 
         # UserMessage(expr) - creates a user message dict
-        match = re.match(r"UserMessage\((.+)\)$", formula.strip())
-        if match:
-            inner_expr = match.group(1).strip()
+        inner_args = whole_call_args(formula, "UserMessage")
+        if inner_args:
+            inner_expr = inner_args.strip()
             # Evaluate the inner expression
             text = self.eval(f"={inner_expr}")
             return {"role": "user", "text": str(text) if text else ""}
 
         # AgentMessage(expr) - creates an assistant message dict
-        match = re.match(r"AgentMessage\((.+)\)$", formula.strip())
-        if match:
-            inner_expr = match.group(1).strip()
+        inner_args = whole_call_args(formula, "AgentMessage")
+        if inner_args:
+            inner_expr = inner_args.strip()
             text = self.eval(f"={inner_expr}")
             return {"role": "assistant", "text": str(text) if text else ""}
 
         # MessageText(expr) - extracts text from the last message
-        match = re.match(r"MessageText\((.+)\)$", formula.strip())
-        if match:
-            inner_expr = match.group(1).strip()
+        inner_args = whole_call_args(formula, "MessageText")
+        if inner_args:
+            inner_expr = inner_args.strip()
             # Reuse the helper method for consistent text extraction
             return self._eval_and_replace_message_text(inner_expr)
 

python/packages/declarative/agent_framework_declarative/_workflows/_powerfx_functions.py

@@ -12,9 +12,156 @@
 
 from __future__ import annotations
 
+from collections.abc import Iterator
 from typing import Any, cast
 
 
+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)
+        # An unterminated /* is not a complete comment: leave it as code so a
+        # trailing one is not mistaken for trivia following a complete call.
+        return start 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
+
+
+def _after_trailing_trivia(formula: str, start: int) -> int:
+    """Return the index after trailing whitespace and complete comments.
+
+    Unlike the expression walker, this never skips a string literal: a quoted
+    token after the call's own closing paren is trailing code, not trivia, so it
+    stops the scan and leaves the formula to PowerFx.
+    """
+    index = start
+    while index < len(formula):
+        if formula[index].isspace():
+            index += 1
+            continue
+        if formula.startswith(("//", "/*"), index):
+            token_end = _skip_powerfx_opaque_token(formula, index)
+            if token_end == index:
+                # Unterminated /* is not a complete comment, so real code follows.
+                break
+            index = token_end
+            continue
+        break
+    return index
+
+
+def whole_call_args(formula: str, *names: str) -> str | None:
+    """Return the argument text when ``formula`` is exactly one ``Name(...)`` call.
+
+    A custom-function handler must not fire for a formula that merely *starts*
+    with its name: the closing parenthesis has to be the call's own and nothing
+    but trivia may follow it. Otherwise a trailing operator is swallowed into the
+    argument list and the formula is evaluated as something the author never
+    wrote, instead of being left to PowerFx.
+
+    Parentheses inside string literals and comments are data, not delimiters, so
+    the scan walks the formula with the shared PowerFx expression walker: it
+    skips those tokens whole and counts only parentheses that are real code,
+    including those inside an interpolated string's ``{...}`` expression sections.
+
+    Args:
+        formula: The stripped formula to inspect.
+        *names: Candidate function names, matched case-sensitively.
+
+    Returns:
+        The text between the call's parentheses - an empty string for an
+        argument-less call - or None when ``formula`` is not a single complete
+        call to one of ``names``. Callers that require an argument check the
+        result for emptiness so the call still falls through to PowerFx, which
+        reports the arity error.
+    """
+    for name in names:
+        prefix = f"{name}("
+        if not formula.startswith(prefix):
+            continue
+        positions = _iter_powerfx_expression_indices(formula)
+        depth = 0
+        close_index = -1
+        for index in positions:
+            char = formula[index]
+            if char == "(":
+                depth += 1
+            elif char == ")":
+                depth -= 1
+                if depth == 0:
+                    close_index = index
+                    break
+        if close_index == -1:
+            # Unbalanced: the call never closes, so it is not a whole call.
+            return None
+        # Only whitespace and complete comments may follow the call's own closing
+        # paren. The suffix scan deliberately does not reuse the literal-skipping
+        # walker: a trailing string literal (or any other code) there means the
+        # formula is a larger expression the caller must not claim.
+        if _after_trailing_trivia(formula, close_index + 1) != len(formula):
+            return None
+        return formula[len(prefix) : close_index]
+    return None
+
+
 def message_text(messages: Any) -> str:
     """Extract text content from a message or list of messages.
 

python/packages/declarative/agent_framework_declarative/_workflows/_state.py

@@ -441,7 +441,7 @@ def _eval_simple(self, formula: str) -> Any:
         Returns:
             The evaluated result
         """
-        from ._powerfx_functions import CUSTOM_FUNCTIONS
+        from ._powerfx_functions import CUSTOM_FUNCTIONS, whole_call_args
 
         formula = formula.strip()
 
@@ -452,15 +452,15 @@ def _eval_simple(self, formula: str) -> Any:
             return not bool(result)
 
         # Handle Not() function
-        if formula.startswith("Not(") and formula.endswith(")"):
-            inner = formula[4:-1].strip()
-            result = self._eval_simple(inner)
+        not_args = whole_call_args(formula, "Not")
+        if not_args is not None:
+            result = self._eval_simple(not_args.strip())
             return not bool(result)
 
         # Handle function calls
         for func_name, func in CUSTOM_FUNCTIONS.items():
-            if formula.startswith(f"{func_name}(") and formula.endswith(")"):
-                args_str = formula[len(func_name) + 1 : -1]
+            args_str = whole_call_args(formula, func_name)
+            if args_str is not None:
                 # Simple argument parsing (doesn't handle nested calls well)
                 args = self._parse_function_args(args_str)
                 evaluated_args = [self._eval_simple(arg) if isinstance(arg, str) else arg for arg in args]