Python: [BREAKING] bound PowerFx state construction

python/packages/core/agent_framework/_workflows/_state.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import copy
+from collections.abc import Callable
 from typing import Any
 
 
@@ -70,6 +71,16 @@ def has(self, key: str) -> bool:
             return self._pending[key] is not _DeleteSentinel
         return key in self._committed
 
+    def _validate(self, key: str, validator: Callable[[Any], None]) -> None:
+        """Validate a stored value before copying it.
+
+        The internal validator must not mutate or retain the value. Missing and
+        pending-deleted keys are not validated, matching ``get`` visibility.
+        """
+        value = self._pending.get(key, self._committed.get(key, _DeleteSentinel))
+        if value is not _DeleteSentinel:
+            validator(value)
+
     def delete(self, key: str) -> None:
         """Mark a key for deletion.
 

python/packages/core/tests/workflow/test_state.py

@@ -2,6 +2,8 @@
 
 """Unit tests for the State class superstep caching behavior."""
 
+from unittest.mock import MagicMock
+
 import pytest
 
 from agent_framework import SecretString
@@ -68,6 +70,41 @@ def test_get_with_default(self) -> None:
         assert state.get("missing") is None
         assert state.get("missing", "default") == "default"
 
+    def test_validation_uses_pending_then_committed_values(self) -> None:
+        state = State()
+        state.set("key", {"value": "committed"})
+        state.commit()
+        validator = MagicMock()
+
+        state._validate("key", validator)
+        validator.assert_called_once_with({"value": "committed"})
+        validator.reset_mock()
+        state.set("key", {"value": "pending"})
+        state._validate("key", validator)
+        validator.assert_called_once_with({"value": "pending"})
+
+    def test_validation_skips_missing_and_pending_deleted_values(self) -> None:
+        state = State()
+        state.set("key", "value")
+        state.commit()
+        state.delete("key")
+        validator = MagicMock()
+
+        state._validate("missing", validator)
+        state._validate("key", validator)
+
+        validator.assert_not_called()
+
+    def test_validation_failure_leaves_state_unchanged(self) -> None:
+        state = State()
+        state.set("key", {"value": "original"})
+        validator = MagicMock(side_effect=ValueError("invalid"))
+
+        with pytest.raises(ValueError, match="invalid"):
+            state._validate("key", validator)
+
+        assert state.get("key") == {"value": "original"}
+
     def test_has_returns_true_for_existing_key(self) -> None:
         state = State()
         state.set("key", "value")

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

@@ -46,6 +46,7 @@
 from agent_framework._workflows._state import State
 
 from ._errors import DeclarativeWorkflowError
+from ._powerfx_limits import _PowerFxStateBudget, _validate_powerfx_state  # pyright: ignore[reportPrivateUsage]
 
 try:
     from powerfx import Engine
@@ -169,6 +170,8 @@ def __post_init__(self) -> None:
         # contents of ``values`` / ``referenced_names``: caller mutations
         # to the original objects after construction cannot leak into
         # ``resolve()``.
+        _validate_powerfx_state(self.values)
+        _validate_powerfx_state(self.referenced_names)
         object.__setattr__(self, "values", MappingProxyType(dict(self.values)))
         object.__setattr__(self, "referenced_names", frozenset(self.referenced_names))
 
@@ -181,6 +184,8 @@ def resolve(self) -> dict[str, str]:
         unrelated environment variables never enter the PowerFx scope.
         Configuration values always win over the environment fallback.
         """
+        _validate_powerfx_state(self.values)
+        _validate_powerfx_state(self.referenced_names)
         resolved = {name: str(value) for name, value in self.values.items()}
         if self.restrict_to_configuration:
             return resolved
@@ -314,10 +319,16 @@ def _make_powerfx_safe(value: Any) -> Any:
 
     Returns:
         A PowerFx-safe representation of the value
+
+    Raises:
+        ValueError: If the input or converted symbols exceed the state budget.
     """
-    if value is None:
-        return value
+    _validate_powerfx_state(value)
+    return _convert_powerfx_value(value, _PowerFxStateBudget(), 0)
 
+
+def _convert_powerfx_value(value: Any, budget: _PowerFxStateBudget, depth: int) -> Any:
+    """Convert an already bounded value without restarting its traversal budget."""
     # Enum coercion must run BEFORE the primitive type check: many MAF
     # enums (e.g. MessageRole) are ``str``-subclass enums, so they pass
     # ``isinstance(v, str)`` but pythonnet refuses to convert them to
@@ -326,25 +337,27 @@ def _make_powerfx_safe(value: Any) -> Any:
     # to the underlying value (or its string form) so PowerFx sees a
     # plain ``str``/``int``.
     if isinstance(value, Enum):
-        return _make_powerfx_safe(value.value)
+        return _convert_powerfx_value(value.value, budget, depth)
+
+    if not isinstance(value, (*_POWERFX_SAFE_TYPES, dict, list)):
+        if hasattr(value, "__dict__"):
+            return _convert_powerfx_value(vars(value), budget, depth)
+        value = str(value)
 
+    budget.consume(value, depth)
     if isinstance(value, _POWERFX_SAFE_TYPES):
         return value
 
     if isinstance(value, dict):
         value_dict = cast(Mapping[Any, Any], value)
-        return {str(k): _make_powerfx_safe(v) for k, v in value_dict.items()}
-
-    if isinstance(value, list):
-        value_list = cast(list[Any], value)
-        return [_make_powerfx_safe(item) for item in value_list]
-
-    # Try to convert objects with __dict__ or dataclass-style attributes
-    if hasattr(value, "__dict__"):
-        return _make_powerfx_safe(vars(value))
+        result: dict[str, Any] = {}
+        for key, member in value_dict.items():
+            name = str(key)
+            budget.consume(name, depth + 1)
+            result[name] = _convert_powerfx_value(member, budget, depth + 1)
+        return result
 
-    # For other objects, try to convert to string representation
-    return str(value)
+    return [_convert_powerfx_value(item, budget, depth + 1) for item in cast(list[Any], value)]
 
 
 class DeclarativeWorkflowState:
@@ -384,6 +397,7 @@ def initialize(self, inputs: Mapping[str, Any] | None = None) -> None:
         Args:
             inputs: Initial workflow inputs (become Workflow.Inputs.*)
         """
+        _validate_powerfx_state(inputs)
         conversation_id = str(uuid.uuid4())
         state_data: DeclarativeStateData = {
             "Inputs": dict(inputs) if inputs else {},
@@ -402,10 +416,11 @@ def initialize(self, inputs: Mapping[str, Any] | None = None) -> None:
             "Conversation": {"messages": [], "history": []},
             "Custom": {},
         }
-        self._state.set(DECLARATIVE_STATE_KEY, state_data)
+        self.set_state_data(state_data)
 
     def get_state_data(self) -> DeclarativeStateData:
         """Get the full state data dict from state."""
+        self._state._validate(DECLARATIVE_STATE_KEY, _validate_powerfx_state)  # pyright: ignore[reportPrivateUsage]
         result = self._state.get(DECLARATIVE_STATE_KEY)
         if result is None:
             # Initialize if not present
@@ -421,10 +436,12 @@ def is_initialized(self) -> bool:
         scenarios), the start executor needs to avoid calling initialize()
         and clobbering the prior turn's Conversation/Local/System data.
         """
+        self._state._validate(DECLARATIVE_STATE_KEY, _validate_powerfx_state)  # pyright: ignore[reportPrivateUsage]
         return self._state.get(DECLARATIVE_STATE_KEY) is not None
 
     def set_state_data(self, data: DeclarativeStateData) -> None:
         """Set the full state data dict in state."""
+        _validate_powerfx_state(data)
         self._state.set(DECLARATIVE_STATE_KEY, data)
 
     def get(self, path: str, default: Any = None) -> Any:
@@ -617,6 +634,8 @@ def eval(self, expression: str) -> Any:
         Raises:
             RuntimeError: If the powerfx package is not installed and the
                 expression requires PowerFx evaluation.
+            ValueError: If state copying or symbol conversion exceeds the
+                PowerFx state budget.
         """
         if not expression:
             return expression
@@ -940,10 +959,10 @@ def _to_powerfx_symbols(self) -> dict[str, Any]:
             symbols["Env"] = env_bound
         # Debug log the Local symbols to help diagnose type issues
         if local_data:
-            for key, value in local_data.items():
+            for value in local_data.values():
                 logger.debug(
-                    f"PowerFx symbol Local.{key}: type={type(value).__name__}, "
-                    f"value_preview={str(value)[:100] if value else None}"
+                    "PowerFx Local symbol type=%s",
+                    type(value).__name__,
                 )
         result = _make_powerfx_safe(symbols)
         return cast(dict[str, Any], result)

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

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Structural budgets for state copying and PowerFx symbol conversion.
+
+Each traversal has its own budget and rejects cycles or excess with ValueError,
+rather than truncating values. These limits do not bound expression execution
+or application-defined Python copy/conversion hooks.
+"""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, cast
+
+_MAX_POWERFX_STATE_DEPTH = 64
+_MAX_POWERFX_STATE_NODES = 10_000
+_MAX_POWERFX_STATE_TEXT_SIZE = 1_048_576
+
+
+class _PowerFxStateLimitError(ValueError):
+    """State cannot be copied or marshalled within the PowerFx budget."""
+
+
+@dataclass
+class _PowerFxStateBudget:
+    """Count values, containers, and mapping keys, including repeated aliases.
+
+    Depth starts at zero. Text size counts string characters and binary bytes,
+    not encoded size or total memory.
+    """
+
+    nodes: int = 0
+    text_size: int = 0
+
+    def consume(self, item: Any, depth: int) -> None:
+        self.nodes += 1
+        if self.nodes > _MAX_POWERFX_STATE_NODES:
+            raise _PowerFxStateLimitError("PowerFx state exceeds the node budget")
+        if depth > _MAX_POWERFX_STATE_DEPTH:
+            raise _PowerFxStateLimitError("PowerFx state exceeds the depth budget")
+        if isinstance(item, (str, bytes, bytearray)):
+            self.text_size += len(item)
+            if self.text_size > _MAX_POWERFX_STATE_TEXT_SIZE:
+                raise _PowerFxStateLimitError("PowerFx state exceeds the text size budget")
+
+
+def _validate_powerfx_state(value: Any) -> None:  # pyright: ignore[reportUnusedFunction]
+    """Reject cyclic or over-budget data before copying or conversion."""
+    budget = _PowerFxStateBudget()
+    active: set[int] = set()
+
+    def visit(item: Any, depth: int) -> None:
+        budget.consume(item, depth)
+        if item is None or isinstance(item, (str, bytes, bytearray, bool, int, float)):
+            return
+
+        identity = id(item)
+        if identity in active:
+            raise _PowerFxStateLimitError("PowerFx state contains a cycle")
+        active.add(identity)
+        try:
+            if isinstance(item, Enum):
+                visit(item.value, depth + 1)
+            elif isinstance(item, Mapping):
+                for key, member in cast(Mapping[Any, Any], item).items():
+                    visit(key, depth + 1)
+                    visit(member, depth + 1)
+            elif isinstance(item, (list, tuple, set, frozenset)):
+                for member in cast(list[Any] | tuple[Any, ...] | set[Any] | frozenset[Any], item):
+                    visit(member, depth + 1)
+            elif hasattr(item, "__dict__"):
+                visit(vars(item), depth + 1)
+        finally:
+            active.remove(identity)
+
+    visit(value, 0)

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

@@ -16,6 +16,8 @@
 from collections.abc import Mapping
 from typing import Any, cast
 
+from ._powerfx_limits import _PowerFxStateLimitError, _validate_powerfx_state  # pyright: ignore[reportPrivateUsage]
+
 try:
     from powerfx import Engine
 
@@ -105,6 +107,7 @@ def __init__(
             inputs: Initial inputs to the workflow. These become available
                    as Workflow.Inputs.* and are immutable after initialization.
         """
+        _validate_powerfx_state(inputs)
         self._inputs: dict[str, Any] = dict(inputs) if inputs else {}
         self._local: dict[str, Any] = {}
         self._outputs: dict[str, Any] = {}
@@ -329,7 +332,19 @@ def to_powerfx_symbols(self) -> dict[str, Any]:
 
         Returns:
             A dictionary suitable for passing to PowerFx Engine.eval()
+
+        Raises:
+            ValueError: If the state or projected symbols exceed the PowerFx state budget.
         """
+        _validate_powerfx_state({
+            "Inputs": self._inputs,
+            "Outputs": self._outputs,
+            "Local": self._local,
+            "System": self._system,
+            "Agent": self._agent,
+            "Conversation": self._conversation,
+            "Custom": self._custom,
+        })
         symbols = {
             "Workflow": {
                 "Inputs": dict(self._inputs),
@@ -345,11 +360,12 @@ def to_powerfx_symbols(self) -> dict[str, Any]:
         }
         # Debug log the Local symbols to help diagnose type issues
         if self._local:
-            for key, value in self._local.items():
+            for value in self._local.values():
                 logger.debug(
-                    f"PowerFx symbol Local.{key}: type={type(value).__name__}, "
-                    f"value_preview={str(value)[:100] if value else None}"
+                    "PowerFx Local symbol type=%s",
+                    type(value).__name__,
                 )
+        _validate_powerfx_state(symbols)
         return symbols
 
     def eval(self, expression: str) -> Any:
@@ -363,6 +379,9 @@ def eval(self, expression: str) -> Any:
 
         Returns:
             The evaluated result, or the original expression if not a PowerFx expression
+
+        Raises:
+            ValueError: If symbol construction exceeds the PowerFx state budget.
         """
         if not expression:
             return expression
@@ -378,6 +397,8 @@ def eval(self, expression: str) -> Any:
             try:
                 symbols = self.to_powerfx_symbols()
                 return _powerfx_engine.eval(formula, symbols=symbols)
+            except _PowerFxStateLimitError:
+                raise
             except Exception as exc:
                 logger.warning(f"PowerFx evaluation failed for '{expression[:50]}': {exc}")
                 # Fall through to simple evaluation