diff --git a/python/packages/core/agent_framework/_workflows/_state.py b/python/packages/core/agent_framework/_workflows/_state.py index 64759da72c..136b70de14 100644 --- a/python/packages/core/agent_framework/_workflows/_state.py +++ b/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. diff --git a/python/packages/core/tests/workflow/test_state.py b/python/packages/core/tests/workflow/test_state.py index e233eb976f..6b6a52c608 100644 --- a/python/packages/core/tests/workflow/test_state.py +++ b/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") 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 402baea186..a9eb95473d 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py +++ b/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) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_powerfx_limits.py b/python/packages/declarative/agent_framework_declarative/_workflows/_powerfx_limits.py new file mode 100644 index 0000000000..96de2610ca --- /dev/null +++ b/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) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_state.py b/python/packages/declarative/agent_framework_declarative/_workflows/_state.py index 9ba9dd964b..73f06230ad 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_state.py +++ b/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 diff --git a/python/packages/declarative/tests/test_powerfx_safe.py b/python/packages/declarative/tests/test_powerfx_safe.py index fccbd72b28..871e14cdb4 100644 --- a/python/packages/declarative/tests/test_powerfx_safe.py +++ b/python/packages/declarative/tests/test_powerfx_safe.py @@ -10,8 +10,18 @@ pin down the Enum coercion branch so we don't regress that interop fix. """ +from dataclasses import dataclass +from decimal import Decimal from enum import Enum, IntEnum +from typing import Any +from unittest.mock import MagicMock +import pytest +from agent_framework._workflows._state import State + +from agent_framework_declarative._workflows import _declarative_base as base +from agent_framework_declarative._workflows import _powerfx_limits as limits +from agent_framework_declarative._workflows import _state as legacy from agent_framework_declarative._workflows._declarative_base import _make_powerfx_safe @@ -57,3 +67,231 @@ def test_enum_inside_list_is_coerced(): assert safe == ["user", 1] assert type(safe[0]) is str assert type(safe[1]) is int + + +@pytest.mark.parametrize("committed", [False, True]) +def test_state_depth_is_checked_before_copy_and_engine_dispatch( + monkeypatch: pytest.MonkeyPatch, committed: bool +) -> None: + """Even an expression without state references must respect the snapshot budget.""" + state = State() + workflow_state = base.DeclarativeWorkflowState(state) + workflow_state.initialize() + data = workflow_state.get_state_data() + nested: Any = "leaf" + for _ in range(65): + nested = [nested] + data["Local"]["unused"] = nested + state.set(base.DECLARATIVE_STATE_KEY, data) + if committed: + state.commit() + + engine = MagicMock() + monkeypatch.setattr(base, "Engine", engine) + deepcopy = MagicMock(side_effect=AssertionError("State was copied before validating its budget")) + monkeypatch.setattr("agent_framework._workflows._state.copy.deepcopy", deepcopy) + + with pytest.raises(ValueError, match="PowerFx state.*depth"): + workflow_state.eval("=1 + 1") + + deepcopy.assert_not_called() + engine.assert_not_called() + + +@pytest.mark.parametrize( + ("setting", "limit", "accepted", "rejected", "reason"), + [ + ("_MAX_POWERFX_STATE_DEPTH", 2, [[0]], [[[0]]], "depth"), + ("_MAX_POWERFX_STATE_NODES", 3, [0, 1], [0, 1, 2], "node"), + ("_MAX_POWERFX_STATE_TEXT_SIZE", 4, {"ab": "cd"}, {"ab": "cde"}, "text size"), + ], +) +def test_budget_boundaries( + monkeypatch: pytest.MonkeyPatch, + setting: str, + limit: int, + accepted: Any, + rejected: Any, + reason: str, +) -> None: + monkeypatch.setattr(limits, setting, limit) + assert _make_powerfx_safe(accepted) == accepted + with pytest.raises(ValueError, match=f"PowerFx state.*{reason}"): + _make_powerfx_safe(rejected) + + +def test_shared_values_count_at_each_emitted_occurrence(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_NODES", 5) + shared = [1] + assert _make_powerfx_safe([shared, shared]) == [[1], [1]] + with pytest.raises(ValueError, match="node budget"): + _make_powerfx_safe([shared, shared, shared]) + + +@pytest.mark.parametrize("kind", ["list", "dict", "object"]) +def test_cycles_are_rejected(kind: str) -> None: + if kind == "list": + value: Any = [] + value.append(value) + elif kind == "dict": + value = {} + value["self"] = value + else: + + @dataclass + class Record: + child: Any = None + + value = Record() + value.child = value + with pytest.raises(ValueError, match="PowerFx state contains a cycle"): + _make_powerfx_safe(value) + + +def test_conversion_preserves_normal_data() -> None: + @dataclass + class Record: + role: _StrRole + score: Decimal + + assert _make_powerfx_safe({"message": Record(_StrRole.USER, Decimal("1.25")), 7: [None, True]}) == { + "message": {"role": "user", "score": Decimal("1.25")}, + "7": [None, True], + } + + +def test_validated_snapshot_preserves_copy_isolation() -> None: + store = State() + state = base.DeclarativeWorkflowState(store) + state.initialize({"question": "hello"}) + state.set("Local.items", [1, 2]) + store.commit() + + snapshot = state.get_state_data() + snapshot["Local"]["items"].append(3) + assert state.get("Local.items") == [1, 2] + state.set_state_data(snapshot) + assert state.get("Local.items") == [1, 2, 3] + store.discard() + assert state.get("Local.items") == [1, 2] + + +def test_budget_resets_between_conversions(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_NODES", 2) + assert _make_powerfx_safe([1]) == [1] + assert _make_powerfx_safe([2]) == [2] + + +def test_configuration_is_checked_before_snapshotting(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_TEXT_SIZE", 3) + with pytest.raises(ValueError, match="text size budget"): + base.DeclarativeEnvConfig(values={"name": "value"}) + + +def test_symbol_logging_does_not_format_state_values() -> None: + class Record: + def __str__(self) -> str: + raise AssertionError("Symbol logging must not format whole values") + + state = base.DeclarativeWorkflowState(State()) + state.initialize() + state.set("Local.record", Record()) + + assert state._to_powerfx_symbols()["Local"]["record"] == {} + + +def test_converted_strings_consume_one_aggregate_budget(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[None] = [] + + class Text: + __slots__ = () + + def __str__(self) -> str: + calls.append(None) + return "text" + + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_TEXT_SIZE", 7) + with pytest.raises(ValueError, match="text size budget"): + _make_powerfx_safe([Text(), Text(), Text()]) + assert len(calls) == 2 + + +def test_state_write_rejection_preserves_previous_value(monkeypatch: pytest.MonkeyPatch) -> None: + state = base.DeclarativeWorkflowState(State()) + state.initialize() + state.set("Local.value", "before") + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_DEPTH", 5) + + with pytest.raises(ValueError, match="depth budget"): + state.set("Local.value", [[[[0]]]]) + + assert state.get("Local.value") == "before" + + +def test_symbol_aliases_share_the_output_budget(monkeypatch: pytest.MonkeyPatch) -> None: + state = base.DeclarativeWorkflowState(State()) + state.initialize({"value": "ordinary input"}) + data = state.get_state_data() + size = limits._PowerFxStateBudget() + # Capture the exact small fixture's text use without depending on generated IDs. + original_consume = limits._PowerFxStateBudget.consume + + def consume(self: Any, item: Any, depth: int) -> None: + original_consume(self, item, depth) + size.text_size = self.text_size + + monkeypatch.setattr(limits._PowerFxStateBudget, "consume", consume) + limits._validate_powerfx_state(data) + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_TEXT_SIZE", size.text_size) + + with pytest.raises(ValueError, match="text size budget"): + state._to_powerfx_symbols() + + +def test_message_text_cleanup_after_symbol_budget_failure(monkeypatch: pytest.MonkeyPatch) -> None: + state = base.DeclarativeWorkflowState(State()) + state.initialize() + state.set("Local._TempMessageText0", "original") + before = state.get_state_data() + monkeypatch.setattr(state, "_eval_and_replace_message_text", lambda expression: "hello") + monkeypatch.setattr(state, "_to_powerfx_symbols", MagicMock(side_effect=limits._PowerFxStateLimitError("budget"))) + monkeypatch.setattr(base, "Engine", MagicMock()) + + with pytest.raises(ValueError, match="budget"): + state.eval("=Upper(MessageText(Local.Messages))") + + assert state.get_state_data() == before + + +def test_rejected_temporary_binding_preserves_state(monkeypatch: pytest.MonkeyPatch) -> None: + state = base.DeclarativeWorkflowState(State()) + state.initialize() + state.set("Local._TempMessageText0", "original") + before = state.get_state_data() + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_TEXT_SIZE", 400) + limits._validate_powerfx_state(before) + monkeypatch.setattr(state, "_eval_and_replace_message_text", lambda expression: "x" * 401) + engine = MagicMock() + monkeypatch.setattr(base, "Engine", engine) + + with pytest.raises(ValueError, match="text size budget"): + state.eval("=Upper(MessageText(Local.Messages))") + + assert state.get_state_data() == before + engine.assert_not_called() + + +def test_legacy_budget_failure_does_not_fall_back(monkeypatch: pytest.MonkeyPatch) -> None: + state = legacy.WorkflowState() + state.set("Local.values", [1, 2, 3]) + engine = MagicMock() + fallback = MagicMock() + monkeypatch.setattr(legacy, "_powerfx_engine", engine) + monkeypatch.setattr(state, "_eval_simple", fallback) + monkeypatch.setattr(limits, "_MAX_POWERFX_STATE_NODES", 3) + + with pytest.raises(ValueError, match="node budget"): + state.eval("=1 + 1") + + engine.eval.assert_not_called() + fallback.assert_not_called()