Python: add powerfx safe mode (#3028)

* add powerfx safe mode

* improved docstring and aligned env_file loading

* ensured test uses reset
This commit is contained in:
Eduard van Valkenburg
2025-12-23 21:12:50 +01:00
committed by GitHub
Unverified
parent 5ab47596ff
commit 4b8a545589
4 changed files with 242 additions and 23 deletions
@@ -37,6 +37,7 @@ from ._models import (
RemoteConnection,
Tool,
WebSearchTool,
_safe_mode_context,
agent_schema_dispatch,
)
@@ -118,7 +119,9 @@ class AgentFactory:
client_kwargs: Mapping[str, Any] | None = None,
additional_mappings: Mapping[str, ProviderTypeMapping] | None = None,
default_provider: str = "AzureAIClient",
env_file: str | None = None,
safe_mode: bool = True,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Create the agent factory, with bindings.
@@ -151,7 +154,15 @@ class AgentFactory:
that accepts the model.id value.
default_provider: The default provider used when model.provider is not specified,
default is "AzureAIClient".
env_file: An optional path to a .env file to load environment variables from.
safe_mode: Whether to run in safe mode, default is True.
When safe_mode is True, environment variables are not accessible in the powerfx expressions.
You can still use environment variables, but through the constructors of the classes.
Which means you must make sure you are using the standard env variable names of the classes
you are using and not custom ones and remove the powerfx statements that start with `=Env.`.
Only when you trust the source of your yaml files, you can set safe_mode to False
via the AgentFactory constructor.
env_file_path: The path to the .env file to load environment variables from.
env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
"""
self.chat_client = chat_client
self.bindings = bindings
@@ -159,7 +170,8 @@ class AgentFactory:
self.client_kwargs = client_kwargs or {}
self.additional_mappings = additional_mappings or {}
self.default_provider: str = default_provider
load_dotenv(dotenv_path=env_file)
self.safe_mode = safe_mode
load_dotenv(dotenv_path=env_file_path, encoding=env_file_encoding)
def create_agent_from_yaml_path(self, yaml_path: str | Path) -> ChatAgent:
"""Create a ChatAgent from a YAML file path.
@@ -215,6 +227,8 @@ class AgentFactory:
ModuleNotFoundError: If the required module for the provider type cannot be imported.
AttributeError: If the required class for the provider type cannot be found in the module.
"""
# Set safe_mode context before parsing YAML to control PowerFx environment variable access
_safe_mode_context.set(self.safe_mode)
prompt_agent = agent_schema_dispatch(yaml.safe_load(yaml_str))
if not isinstance(prompt_agent, PromptAgent):
raise DeclarativeLoaderError("Only yaml definitions for a PromptAgent are supported for agent creation.")
@@ -2,6 +2,7 @@
import os
import sys
from collections.abc import MutableMapping
from contextvars import ContextVar
from typing import Any, Literal, TypeVar, Union
from agent_framework import get_logger
@@ -21,6 +22,11 @@ else:
logger = get_logger("agent_framework.declarative")
# Context variable for safe_mode setting.
# When True (default), environment variables are NOT accessible in PowerFx expressions.
# When False, environment variables CAN be accessed via Env symbol in PowerFx.
_safe_mode_context: ContextVar[bool] = ContextVar("safe_mode", default=True)
@overload
def _try_powerfx_eval(value: None, log_value: bool = True) -> None: ...
@@ -49,6 +55,9 @@ def _try_powerfx_eval(value: str | None, log_value: bool = True) -> str | None:
)
return value
try:
safe_mode = _safe_mode_context.get()
if safe_mode:
return engine.eval(value[1:])
return engine.eval(value[1:], symbols={"Env": dict(os.environ)})
except Exception as exc:
if log_value:
@@ -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)
@@ -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: