From c30f104b207ea28cade06233a4dd55faf722e9bd Mon Sep 17 00:00:00 2001 From: Peter Ibekwe Date: Mon, 18 May 2026 14:43:26 -0700 Subject: [PATCH] Python: address PR review on declarative toolbox sample Two security fixes for PR #5933: 1. Add safe_mode flag to WorkflowFactory (default True) mirroring AgentFactory. Gates =Env.* exposure inside DeclarativeWorkflowState PowerFx symbols via _safe_mode_context, so workflow YAML loaded from untrusted sources no longer leaks the host's full os.environ snapshot into PowerFx evaluation. The flag is also forwarded to the internally-constructed AgentFactory so inline agent definitions follow the same policy. 2. Pin the invoke_foundry_toolbox_mcp sample's _client_provider to the resolved toolbox endpoint. The bearer-authenticated httpx client is now only returned when MCPToolInvocation.server_url matches the toolbox URL case-insensitively; any other URL gets None (the default unauthenticated path), preventing the Foundry AAD bearer token from being attached to a mis-configured or injected server URL. Mirrors the .NET sample's httpClientProvider guard. The sample is updated to opt in to safe_mode=False because its YAML intentionally uses =Env.FOUNDRY_TOOLBOX_* to keep configuration in env vars under the developer's control. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../_workflows/_declarative_base.py | 11 ++++++-- .../_workflows/_factory.py | 23 +++++++++++++++- .../invoke_foundry_toolbox_mcp/main.py | 27 ++++++++++++++++++- 3 files changed, 57 insertions(+), 4 deletions(-) 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 68105186b5..eac40ea978 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py @@ -43,6 +43,8 @@ from agent_framework import ( ) from agent_framework._workflows._state import State +from .._models import _safe_mode_context + try: from powerfx import Engine except (ImportError, RuntimeError): @@ -710,13 +712,18 @@ class DeclarativeWorkflowState: "Agent": agent_data, "Conversation": conversation_data, "System": system_data, - # Match .NET declarative workflows that use =Env.VAR_NAME. - "Env": dict(os.environ), # Also expose inputs at top level for backward compatibility with =inputs.X syntax "inputs": inputs_data, # Custom namespaces **state_data.get("Custom", {}), } + # Expose ``Env`` ONLY when the active workflow factory has explicitly + # opted out of safe mode. Matches the policy enforced by + # ``_try_powerfx_eval`` in ``_models.py`` and the AgentFactory + # ``safe_mode`` flag. Treat the default as safe even when no factory + # is in scope (e.g. in unit tests) so opt-in is required. + if not _safe_mode_context.get(): + symbols["Env"] = dict(os.environ) # Debug log the Local symbols to help diagnose type issues if local_data: for key, value in local_data.items(): diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index 221dfec3cc..5e907ac5b2 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py @@ -26,6 +26,7 @@ from agent_framework import ( ) from .._loader import AgentFactory +from .._models import _safe_mode_context from ._declarative_builder import DeclarativeWorkflowBuilder from ._errors import DeclarativeWorkflowError from ._http_handler import HttpRequestHandler @@ -93,6 +94,7 @@ class WorkflowFactory: max_iterations: int | None = None, http_request_handler: HttpRequestHandler | None = None, mcp_tool_handler: MCPToolHandler | None = None, + safe_mode: bool = True, ) -> None: """Initialize the workflow factory. @@ -119,6 +121,15 @@ class WorkflowFactory: for a default backed by :class:`agent_framework.MCPStreamableHTTPTool`, or supply your own implementation to enforce SSRF guards, allowlisting, or auth/connection resolution. + safe_mode: Whether to run in safe mode, default is True. + When safe_mode is True, environment variables are NOT accessible from + PowerFx expressions in the workflow YAML (e.g. ``=Env.SOME_VAR`` will + fail to resolve and the original expression string is preserved). + When safe_mode is False, the full ``os.environ`` snapshot is exposed + via the ``Env`` symbol in every PowerFx evaluation. Set safe_mode to + False ONLY when you fully trust the YAML source. The flag is also + forwarded to the internally-constructed :class:`AgentFactory` so + inline agent definitions follow the same policy. Examples: .. code-block:: python @@ -152,7 +163,8 @@ class WorkflowFactory: env_file=".env", ) """ - self._agent_factory = agent_factory or AgentFactory(env_file_path=env_file) + self.safe_mode = safe_mode + self._agent_factory = agent_factory or AgentFactory(env_file_path=env_file, safe_mode=safe_mode) self._agents: dict[str, SupportsAgentRun | AgentExecutor] = dict(agents) if agents else {} self._bindings: dict[str, Any] = dict(bindings) if bindings else {} self._tools: dict[str, Any] = {} # Tool registry for InvokeFunctionTool actions @@ -338,6 +350,15 @@ class WorkflowFactory: # Validate the workflow definition self._validate_workflow_def(workflow_def) + # Set safe_mode context before evaluating any PowerFx expressions. The + # contextvar gates ``Env`` exposure inside ``DeclarativeWorkflowState`` + # symbols and inside ``_try_powerfx_eval``; both check + # ``_safe_mode_context.get()`` at evaluation time. Because the + # contextvar propagates through asyncio tasks spawned from the current + # context, this value persists into ``workflow.run(...)`` invocations + # made on the same coroutine. + _safe_mode_context.set(self.safe_mode) + # Extract workflow metadata # Support both "name" field and trigger.id for workflow name name: str = workflow_def.get("name", "") diff --git a/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/main.py b/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/main.py index 36c799fc91..ceffb08393 100644 --- a/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/main.py +++ b/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/main.py @@ -93,6 +93,7 @@ import asyncio import contextlib import json import os +import sys from collections.abc import Iterator from pathlib import Path from typing import Any @@ -445,7 +446,23 @@ async def main() -> None: follow_redirects=True, ) - async def _client_provider(_inv: MCPToolInvocation) -> httpx.AsyncClient | None: + async def _client_provider(invocation: MCPToolInvocation) -> httpx.AsyncClient | None: + # Pin the bearer-authenticated client to the resolved toolbox URL. + # The Foundry AAD bearer token is scoped to ``https://ai.azure.com`` + # but we still refuse to attach it to any URL we did not provision — + # if the YAML resolves a different ``serverUrl`` (e.g. via a tampered + # ``Env.*`` value or a config injection), returning ``None`` causes + # ``DefaultMCPToolHandler`` to fall back to an unauthenticated client, + # which will fail to authenticate to the proxy instead of forwarding + # the token outbound. Mirrors the .NET sample's + # ``httpClientProvider`` URL guard. + if invocation.server_url.casefold() != toolbox_endpoint.casefold(): + print( + f"[security] Refusing to attach Foundry bearer token to unexpected MCP URL: " + f"{invocation.server_url}", + file=sys.stderr, + ) + return None return http_client async with ( @@ -456,6 +473,14 @@ async def main() -> None: factory = WorkflowFactory( agents={AGENT_NAME: summary_agent}, mcp_tool_handler=mcp_handler, + # The workflow YAML references ``=Env.FOUNDRY_TOOLBOX_*`` to keep + # the sample's toolbox URL / tool names configurable without + # editing the YAML. ``WorkflowFactory`` defaults to ``safe_mode=True`` + # which would block those expressions; this sample opts in to the + # less-safe mode because we control both the YAML and the env + # vars. Do NOT copy this flag into a workflow that loads YAML + # from untrusted sources. + safe_mode=False, ) workflow_path = Path(__file__).parent / "workflow.yaml"