Python: add powerfx safe mode (#3028)

* add powerfx safe mode

* improved docstring and aligned env_file loading

* ensured test uses reset

python/packages/declarative/tests/test_declarative_loader.py

@@ -454,3 +454,140 @@ def test_agent_schema_dispatch_agent_samples(yaml_file: Path, agent_samples_dir:
     result = agent_schema_dispatch(yaml.safe_load(content))
     # Result can be None for unknown kinds, but should not raise exceptions
     assert result is not None, f"agent_schema_dispatch returned None for {yaml_file.relative_to(agent_samples_dir)}"
+
+
+class TestAgentFactorySafeMode:
+    """Tests for AgentFactory safe_mode parameter."""
+
+    def test_agent_factory_safe_mode_default_is_true(self):
+        """Test that safe_mode is True by default."""
+        from agent_framework_declarative._loader import AgentFactory
+
+        factory = AgentFactory()
+        assert factory.safe_mode is True
+
+    def test_agent_factory_safe_mode_can_be_set_false(self):
+        """Test that safe_mode can be explicitly set to False."""
+        from agent_framework_declarative._loader import AgentFactory
+
+        factory = AgentFactory(safe_mode=False)
+        assert factory.safe_mode is False
+
+    def test_agent_factory_safe_mode_blocks_env_in_yaml(self, monkeypatch):
+        """Test that safe_mode=True blocks environment variable access in YAML parsing."""
+        from unittest.mock import MagicMock
+
+        from agent_framework_declarative._loader import AgentFactory
+
+        monkeypatch.setenv("TEST_MODEL_ID", "gpt-4-from-env")
+
+        # Create a mock chat client to avoid needing real provider
+        mock_client = MagicMock()
+
+        yaml_content = """
+kind: Prompt
+name: test-agent
+description: =Env.TEST_DESCRIPTION
+instructions: Hello world
+"""
+        monkeypatch.setenv("TEST_DESCRIPTION", "Description from env")
+
+        # With safe_mode=True (default), Env access should fail and return original value
+        factory = AgentFactory(chat_client=mock_client, safe_mode=True)
+        agent = factory.create_agent_from_yaml(yaml_content)
+
+        # The description should NOT be resolved from env (PowerFx fails, returns original)
+        assert agent.description == "=Env.TEST_DESCRIPTION"
+
+    def test_agent_factory_safe_mode_false_allows_env_in_yaml(self, monkeypatch):
+        """Test that safe_mode=False allows environment variable access in YAML parsing."""
+        from unittest.mock import MagicMock
+
+        from agent_framework_declarative._loader import AgentFactory
+
+        monkeypatch.setenv("TEST_DESCRIPTION", "Description from env")
+
+        # Create a mock chat client to avoid needing real provider
+        mock_client = MagicMock()
+
+        yaml_content = """
+kind: Prompt
+name: test-agent
+description: =Env.TEST_DESCRIPTION
+instructions: Hello world
+"""
+
+        # With safe_mode=False, Env access should work
+        factory = AgentFactory(chat_client=mock_client, safe_mode=False)
+        agent = factory.create_agent_from_yaml(yaml_content)
+
+        # The description should be resolved from env
+        assert agent.description == "Description from env"
+
+    def test_agent_factory_safe_mode_with_api_key_connection(self, monkeypatch):
+        """Test safe_mode with API key connection containing env variable."""
+        from agent_framework_declarative._models import _safe_mode_context
+
+        monkeypatch.setenv("MY_API_KEY", "secret-key-123")
+
+        yaml_content = """
+kind: Prompt
+name: test-agent
+description: Test agent
+instructions: Hello
+model:
+  id: gpt-4
+  provider: OpenAI
+  apiType: Chat
+  connection:
+    kind: key
+    apiKey: =Env.MY_API_KEY
+"""
+
+        # Manually trigger the YAML parsing to check the context is set correctly
+        import yaml as yaml_module
+
+        from agent_framework_declarative._models import agent_schema_dispatch
+
+        token = _safe_mode_context.set(True)  # Ensure we're in safe mode
+        try:
+            result = agent_schema_dispatch(yaml_module.safe_load(yaml_content))
+
+            # The API key should NOT be resolved (still has the PowerFx expression)
+            assert result.model.connection.apiKey == "=Env.MY_API_KEY"
+        finally:
+            _safe_mode_context.reset(token)
+
+    def test_agent_factory_safe_mode_false_resolves_api_key(self, monkeypatch):
+        """Test safe_mode=False resolves API key from environment."""
+        from agent_framework_declarative._models import _safe_mode_context
+
+        monkeypatch.setenv("MY_API_KEY", "secret-key-123")
+
+        yaml_content = """
+kind: Prompt
+name: test-agent
+description: Test agent
+instructions: Hello
+model:
+  id: gpt-4
+  provider: OpenAI
+  apiType: Chat
+  connection:
+    kind: key
+    apiKey: =Env.MY_API_KEY
+"""
+
+        # With safe_mode=False, the API key should be resolved
+        import yaml as yaml_module
+
+        from agent_framework_declarative._models import agent_schema_dispatch
+
+        token = _safe_mode_context.set(False)  # Disable safe mode
+        try:
+            result = agent_schema_dispatch(yaml_module.safe_load(yaml_content))
+
+            # The API key should be resolved from environment
+            assert result.model.connection.apiKey == "secret-key-123"
+        finally:
+            _safe_mode_context.reset(token)

python/packages/declarative/tests/test_declarative_models.py

@@ -41,6 +41,7 @@ from agent_framework_declarative._models import (
     Template,
     ToolResource,
     WebSearchTool,
+    _safe_mode_context,
     _try_powerfx_eval,
 )
 
@@ -874,35 +875,50 @@ class TestTryPowerfxEval:
         monkeypatch.setenv("API_KEY", "secret123")
         monkeypatch.setenv("PORT", "8080")
 
-        # Test basic env access
-        assert _try_powerfx_eval("=Env.TEST_VAR") == "test_value"
-        assert _try_powerfx_eval("=Env.API_KEY") == "secret123"
-        assert _try_powerfx_eval("=Env.PORT") == "8080"
+        # Set safe_mode=False to allow environment variable access
+        token = _safe_mode_context.set(False)
+        try:
+            # Test basic env access
+            assert _try_powerfx_eval("=Env.TEST_VAR") == "test_value"
+            assert _try_powerfx_eval("=Env.API_KEY") == "secret123"
+            assert _try_powerfx_eval("=Env.PORT") == "8080"
+        finally:
+            _safe_mode_context.reset(token)
 
     def test_env_variable_with_string_concatenation(self, monkeypatch):
         """Test env variables with string concatenation operator."""
         monkeypatch.setenv("BASE_URL", "https://api.example.com")
         monkeypatch.setenv("API_VERSION", "v1")
 
-        # Test concatenation with &
-        result = _try_powerfx_eval('=Env.BASE_URL & "/" & Env.API_VERSION')
-        assert result == "https://api.example.com/v1"
+        # Set safe_mode=False to allow environment variable access
+        token = _safe_mode_context.set(False)
+        try:
+            # Test concatenation with &
+            result = _try_powerfx_eval('=Env.BASE_URL & "/" & Env.API_VERSION')
+            assert result == "https://api.example.com/v1"
 
-        # Test concatenation with literals
-        result = _try_powerfx_eval('="API Key: " & Env.API_VERSION')
-        assert result == "API Key: v1"
+            # Test concatenation with literals
+            result = _try_powerfx_eval('="API Key: " & Env.API_VERSION')
+            assert result == "API Key: v1"
+        finally:
+            _safe_mode_context.reset(token)
 
     def test_string_comparison_operators(self, monkeypatch):
         """Test PowerFx string comparison operators."""
         monkeypatch.setenv("ENV_MODE", "production")
 
-        # Equal to - returns bool
-        assert _try_powerfx_eval('=Env.ENV_MODE = "production"') is True
-        assert _try_powerfx_eval('=Env.ENV_MODE = "development"') is False
+        # Set safe_mode=False to allow environment variable access
+        token = _safe_mode_context.set(False)
+        try:
+            # Equal to - returns bool
+            assert _try_powerfx_eval('=Env.ENV_MODE = "production"') is True
+            assert _try_powerfx_eval('=Env.ENV_MODE = "development"') is False
 
-        # Not equal to - returns bool
-        assert _try_powerfx_eval('=Env.ENV_MODE <> "development"') is True
-        assert _try_powerfx_eval('=Env.ENV_MODE <> "production"') is False
+            # Not equal to - returns bool
+            assert _try_powerfx_eval('=Env.ENV_MODE <> "development"') is True
+            assert _try_powerfx_eval('=Env.ENV_MODE <> "production"') is False
+        finally:
+            _safe_mode_context.reset(token)
 
     def test_string_in_operator(self):
         """Test PowerFx 'in' operator for substring testing (case-insensitive)."""
@@ -958,11 +974,54 @@ class TestTryPowerfxEval:
         monkeypatch.setenv("URL_WITH_QUERY", "https://example.com?param=value")
         monkeypatch.setenv("PATH_WITH_SPACES", "C:\\Program Files\\App")
 
-        result = _try_powerfx_eval("=Env.URL_WITH_QUERY")
-        assert result == "https://example.com?param=value"
+        # Set safe_mode=False to allow environment variable access
+        token = _safe_mode_context.set(False)
+        try:
+            result = _try_powerfx_eval("=Env.URL_WITH_QUERY")
+            assert result == "https://example.com?param=value"
 
-        result = _try_powerfx_eval("=Env.PATH_WITH_SPACES")
-        assert result == "C:\\Program Files\\App"
+            result = _try_powerfx_eval("=Env.PATH_WITH_SPACES")
+            assert result == "C:\\Program Files\\App"
+        finally:
+            _safe_mode_context.reset(token)
+
+    def test_safe_mode_blocks_env_access(self, monkeypatch):
+        """Test that safe_mode=True (default) blocks environment variable access."""
+        monkeypatch.setenv("SECRET_VAR", "secret_value")
+
+        # Set safe_mode=True (default)
+        token = _safe_mode_context.set(True)
+        try:
+            # When safe_mode=True, Env is not available and the expression fails,
+            # returning the original value
+            result = _try_powerfx_eval("=Env.SECRET_VAR")
+            assert result == "=Env.SECRET_VAR"
+        finally:
+            _safe_mode_context.reset(token)
+
+    def test_safe_mode_context_isolation(self, monkeypatch):
+        """Test that safe_mode context variable properly isolates env access."""
+        monkeypatch.setenv("TEST_VAR", "test_value")
+
+        # First, set safe_mode=True - should NOT allow env access
+        token = _safe_mode_context.set(True)
+        try:
+            result_safe = _try_powerfx_eval("=Env.TEST_VAR")
+            assert result_safe == "=Env.TEST_VAR"
+
+            # Then, set safe_mode=False - should allow env access
+            token2 = _safe_mode_context.set(False)
+            try:
+                result_unsafe = _try_powerfx_eval("=Env.TEST_VAR")
+                assert result_unsafe == "test_value"
+            finally:
+                _safe_mode_context.reset(token2)
+
+            # After reset, should block again
+            result_safe_again = _try_powerfx_eval("=Env.TEST_VAR")
+            assert result_safe_again == "=Env.TEST_VAR"
+        finally:
+            _safe_mode_context.reset(token)
 
 
 class TestAgentManifest: