Python: avoid repeated workflow reference traversal

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

@@ -45,6 +45,8 @@
 )
 from agent_framework._workflows._state import State
 
+from ._errors import DeclarativeWorkflowError
+
 try:
     from powerfx import Engine
 except (ImportError, RuntimeError):
@@ -197,6 +199,10 @@ def discover_env_references(node: Any) -> set[str]:
     happen to mention ``Env.SOMETHING`` as plain text, the scan only inspects
     strings that begin with ``=`` (PowerFx expression marker, matching the
     convention enforced by :meth:`DeclarativeWorkflowState.eval`).
+    Shared containers are scanned once by identity without Python recursion.
+    Cyclic mappings and lists are rejected rather than silently skipped.
+    This avoids repeated container traversal, but does not impose a definition-size,
+    depth, parsing-time, or execution budget.
 
     Args:
         node: A parsed workflow definition (typically the dict produced by
@@ -205,23 +211,43 @@ def discover_env_references(node: Any) -> set[str]:
     Returns:
         The set of ``Env`` identifier names referenced in PowerFx
         expressions inside ``node``.
+
+    Raises:
+        DeclarativeWorkflowError: If the definition contains a mapping/list cycle.
     """
     names: set[str] = set()
+    active: set[int] = set()
+    # Retain containers so their identities cannot be reused during the walk.
+    visited: dict[int, Mapping[Any, Any] | list[Any]] = {}
+    stack: list[tuple[int | None, Iterator[Any]]] = [(None, iter((node,)))]
 
-    def visit(value: Any) -> None:
+    while stack:
+        parent_id, children = stack[-1]
+        try:
+            value = next(children)
+        except StopIteration:
+            stack.pop()
+            if parent_id is not None:
+                active.remove(parent_id)
+            continue
         if isinstance(value, str):
             if value.startswith("="):
                 names.update(_ENV_REFERENCE_RE.findall(value))
-            return
-        if isinstance(value, Mapping):
-            for inner in cast(Mapping[Any, Any], value).values():
-                visit(inner)
-            return
-        if isinstance(value, list):
-            for item in cast(list[Any], value):
-                visit(item)
+            continue
+        if not isinstance(value, (Mapping, list)):
+            continue
+
+        container = cast(Mapping[Any, Any] | list[Any], value)
+        container_id = id(container)
+        if container_id in active:
+            raise DeclarativeWorkflowError("Cyclic mappings or lists are not supported in workflow definitions.")
+        if container_id in visited:
+            continue
+        visited[container_id] = container
+        active.add(container_id)
+        children = iter(container.values()) if isinstance(container, Mapping) else iter(container)
+        stack.append((container_id, children))
 
-    visit(node)
     return names
 
 

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

@@ -210,7 +210,8 @@ def create_workflow_from_yaml_path(
             An executable Workflow object with action nodes for each YAML action.
 
         Raises:
-            DeclarativeWorkflowError: If the YAML is invalid or cannot be parsed.
+            DeclarativeWorkflowError: If the YAML is invalid, cannot be parsed, or
+                environment-reference discovery encounters a mapping/list cycle.
             FileNotFoundError: If the YAML file doesn't exist.
 
         Examples:
@@ -262,7 +263,8 @@ def create_workflow_from_yaml(
             An executable Workflow object with action nodes for each YAML action.
 
         Raises:
-            DeclarativeWorkflowError: If the YAML is invalid or cannot be parsed.
+            DeclarativeWorkflowError: If the YAML is invalid, cannot be parsed, or
+                environment-reference discovery encounters a mapping/list cycle.
 
         Examples:
             .. code-block:: python
@@ -338,7 +340,8 @@ def create_workflow_from_definition(
             An executable Workflow object with action nodes for each YAML action.
 
         Raises:
-            DeclarativeWorkflowError: If the definition is invalid or missing required fields.
+            DeclarativeWorkflowError: If the definition is invalid, is missing required
+                fields, or environment-reference discovery encounters a mapping/list cycle.
 
         Examples:
             .. code-block:: python

python/packages/declarative/tests/test_workflow_factory.py

@@ -2,14 +2,23 @@
 
 """Unit tests for WorkflowFactory."""
 
+from collections import UserDict
+from collections.abc import Iterator, ValuesView
+from copy import deepcopy
 from pathlib import Path
+from types import MappingProxyType
 from typing import Any, cast
 from unittest.mock import patch
 
 import pytest
+import yaml
 from agent_framework import Message
 
 from agent_framework_declarative._feature_usage import FeatureIndex
+from agent_framework_declarative._workflows._declarative_base import (
+    DeclarativeActionExecutor,
+    discover_env_references,
+)
 from agent_framework_declarative._workflows._errors import DeclarativeWorkflowError
 from agent_framework_declarative._workflows._factory import WorkflowFactory
 
@@ -89,6 +98,116 @@ def test_valid_workflow_marks_declarative_workflow_used(self):
         mark_feature_used.assert_called_once_with(FeatureIndex.DECLARATIVE_WORKFLOW)
 
 
+class TestWorkflowEnvironmentDiscovery:
+    """Environment discovery preserves shared definitions without expanding them repeatedly."""
+
+    def test_shared_containers_are_scanned_once(self) -> None:
+        class CountingMapping(UserDict[str, Any]):
+            visits = 0
+
+            def values(self) -> ValuesView[Any]:
+                self.visits += 1
+                return super().values()
+
+        class CountingList(list[Any]):
+            visits = 0
+
+            def __iter__(self) -> Iterator[Any]:
+                self.visits += 1
+                return super().__iter__()
+
+        leaf = CountingMapping({"expression": "=Env.SHARED & Env.SHARED"})
+        shared_list = CountingList([leaf])
+        shared = [shared_list, shared_list, leaf, {"expression": "=Env.OTHER"}]
+        tree = [[dict(leaf)], [dict(leaf)], dict(leaf), {"expression": "=Env.OTHER"}]
+
+        assert discover_env_references(shared) == discover_env_references(tree) == {"SHARED", "OTHER"}
+        assert leaf.visits == 1
+        assert shared_list.visits == 1
+        assert shared == tree
+        assert shared[0] is shared[1]
+        assert shared_list[0] is leaf
+
+    def test_only_expression_values_contribute_names(self) -> None:
+        definition = {
+            "=Env.KEY": "Env.PLAIN",
+            "nested": [MappingProxyType({"value": "=Env.FIRST & Env.SECOND"}), {"value": "=Env.FIRST"}],
+            "ignored": [None, 7, False, " =Env.LEADING", "Env.TEXT"],
+        }
+
+        assert discover_env_references(definition) == {"FIRST", "SECOND"}
+
+    @pytest.mark.parametrize("source", ["definition", "yaml"])
+    @pytest.mark.parametrize("cycle_kind", ["mapping", "list"])
+    def test_cycles_raise_definition_error(self, source: str, cycle_kind: str) -> None:
+        metadata: dict[str, Any] = {"value": "not-for-error-output"}
+        if cycle_kind == "mapping":
+            metadata["child"] = {"parent": metadata}
+        else:
+            child: list[Any] = []
+            child.append(child)
+            metadata["child"] = child
+        definition = {
+            "name": "cycle-check",
+            "actions": [{"kind": "SendActivity", "activity": "done"}],
+            "metadata": metadata,
+        }
+        factory = WorkflowFactory()
+
+        with pytest.raises(DeclarativeWorkflowError) as exc_info:
+            if source == "yaml":
+                factory.create_workflow_from_yaml(yaml.safe_dump(definition))
+            else:
+                factory.create_workflow_from_definition(definition)
+
+        assert str(exc_info.value) == "Cyclic mappings or lists are not supported in workflow definitions."
+        assert metadata["value"] == "not-for-error-output"
+        if cycle_kind == "mapping":
+            mapping_child = metadata["child"]
+            assert isinstance(mapping_child, dict)
+            assert mapping_child["parent"] is metadata
+        else:
+            list_child = metadata["child"]
+            assert isinstance(list_child, list)
+            assert list_child[0] is list_child
+
+    @pytest.mark.parametrize("source", ["definition", "yaml"])
+    @pytest.mark.parametrize("restrict_env", [True, False])
+    def test_shared_aliases_preserve_env_configuration(
+        self, source: str, restrict_env: bool, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("DISCOVERY_CONFIG", "environment")
+        monkeypatch.setenv("DISCOVERY_FALLBACK", "fallback")
+        monkeypatch.setenv("DISCOVERY_UNREFERENCED", "not-exposed")
+        shared = {"value": "=Env.DISCOVERY_CONFIG & Env.DISCOVERY_FALLBACK"}
+        definition: dict[str, Any] = {
+            "name": "shared-references",
+            "trigger": {"actions": [{"kind": "SendActivity", "activity": "done"}]},
+            "metadata": [shared, shared],
+        }
+        original = deepcopy(definition)
+        factory = WorkflowFactory(
+            configuration={"DISCOVERY_CONFIG": "configured"},
+            restrict_env_to_configuration=restrict_env,
+        )
+        if source == "yaml":
+            workflow = factory.create_workflow_from_yaml(yaml.safe_dump(definition))
+        else:
+            workflow = factory.create_workflow_from_definition(definition)
+
+        executor = workflow.get_start_executor()
+        assert isinstance(executor, DeclarativeActionExecutor)
+        config = executor._declarative_env_config
+        assert config.referenced_names == {"DISCOVERY_CONFIG", "DISCOVERY_FALLBACK"}
+        expected = {"DISCOVERY_CONFIG": "configured"}
+        if not restrict_env:
+            expected["DISCOVERY_FALLBACK"] = "fallback"
+        assert config.resolve() == expected
+        assert definition == original
+        assert definition["metadata"][0] is definition["metadata"][1]
+        assert shared == original["metadata"][0]
+
+
 class TestWorkflowFactoryMessageInput:
     """Tests for declarative workflows started with a single Message."""