Addressed PR comments.

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

@@ -28,12 +28,14 @@ from __future__ import annotations
 import locale
 import logging
 import os
+import re
 import sys
 import uuid
 from collections.abc import Mapping
 from dataclasses import dataclass
 from decimal import Decimal as _Decimal
 from enum import Enum
+from types import MappingProxyType
 from typing import Any, Literal, cast
 
 from agent_framework import (
@@ -43,8 +45,6 @@ from agent_framework import (
 )
 from agent_framework._workflows._state import State
 
-from .._models import _safe_mode_context  # type: ignore[reportPrivateUsage]
-
 try:
     from powerfx import Engine
 except (ImportError, RuntimeError):
@@ -61,6 +61,82 @@ else:
 logger = logging.getLogger(__name__)
 
 
+_ENV_REFERENCE_RE = re.compile(r"\bEnv\.([A-Za-z_][A-Za-z0-9_]*)")
+
+_EMPTY_ENV_VALUES: Mapping[str, str] = MappingProxyType({})
+_EMPTY_ENV_REFERENCES: frozenset[str] = frozenset()
+
+
+@dataclass(frozen=True)
+class DeclarativeEnvConfig:
+    """Configuration that populates the PowerFx ``Env`` symbol for a workflow.
+
+    Mirrors the .NET ``IConfiguration``-driven environment binding in
+    ``Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowDiagnostics.cs``.
+    The configuration values are always exposed under ``Env.<name>`` (matching
+    ``IConfiguration`` semantics), and ``os.environ`` is consulted only when
+    ``restrict_to_configuration`` is ``False`` and the YAML literally
+    references the name in a PowerFx expression.
+
+    Attributes:
+        values: Caller-supplied configuration values resolved by name when
+            the workflow YAML references ``=Env.NAME``. Always exposed in
+            the ``Env`` symbol regardless of ``restrict_to_configuration``.
+        restrict_to_configuration: When ``True`` (default), the ``Env``
+            symbol is populated exclusively from ``values``; ``os.environ``
+            is never consulted. Set to ``False`` to additionally fall back
+            to ``os.environ`` for names absent from ``values`` that the
+            workflow YAML explicitly references.
+        referenced_names: The set of ``Env.NAME`` symbols discovered in
+            PowerFx expressions inside the workflow definition. The
+            ``os.environ`` fallback is constrained to this allowlist so
+            unrelated environment variables never enter the PowerFx scope.
+    """
+
+    values: Mapping[str, str] = _EMPTY_ENV_VALUES
+    restrict_to_configuration: bool = True
+    referenced_names: frozenset[str] = _EMPTY_ENV_REFERENCES
+
+
+_EMPTY_ENV_CONFIG = DeclarativeEnvConfig()
+
+
+def discover_env_references(node: Any) -> set[str]:
+    """Discover ``Env.NAME`` references in PowerFx expressions inside ``node``.
+
+    Walks any nested ``Mapping``/``list``/scalar structure and inspects every
+    string value. To avoid false positives from doc/description fields that
+    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`).
+
+    Args:
+        node: A parsed workflow definition (typically the dict produced by
+            ``yaml.safe_load``).
+
+    Returns:
+        The set of ``Env`` identifier names referenced in PowerFx
+        expressions inside ``node``.
+    """
+    names: set[str] = set()
+
+    def visit(value: Any) -> None:
+        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)
+
+    visit(node)
+    return names
+
+
 class ConversationData(TypedDict):
     """Structure for conversation-related state data.
 
@@ -172,13 +248,18 @@ class DeclarativeWorkflowState:
     - Conversation: Conversation history
     """
 
-    def __init__(self, state: State):
+    def __init__(self, state: State, env_config: DeclarativeEnvConfig = _EMPTY_ENV_CONFIG):
         """Initialize with a State instance.
 
         Args:
             state: The workflow's state for persistence
+            env_config: Configuration that populates the PowerFx ``Env``
+                symbol when ``_to_powerfx_symbols`` is called. Defaults to
+                an empty configuration which results in no ``Env`` binding,
+                matching the safe default of the :class:`WorkflowFactory`.
         """
         self._state = state
+        self._env_config = env_config
 
     def initialize(self, inputs: Mapping[str, Any] | None = None) -> None:
         """Initialize the declarative state with inputs.
@@ -717,13 +798,27 @@ class DeclarativeWorkflowState:
             # 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)
+        # Populate the ``Env`` symbol from the workflow-level
+        # :class:`DeclarativeEnvConfig`. Caller-supplied configuration
+        # values are always exposed (matching .NET ``IConfiguration``
+        # semantics); ``os.environ`` is consulted only when
+        # ``restrict_to_configuration`` is ``False`` and the YAML explicitly
+        # references the name in a PowerFx expression. When both sources
+        # produce no values the ``Env`` symbol is omitted entirely so
+        # ``=Env.X`` resolves to the literal expression string (preserving
+        # the legacy "unbound identifier" fallback behaviour).
+        env_bound: dict[str, str] = {}
+        for name, value in self._env_config.values.items():
+            env_bound[name] = str(value)
+        if not self._env_config.restrict_to_configuration:
+            for name in self._env_config.referenced_names:
+                if name in env_bound:
+                    continue
+                env_value = os.environ.get(name)
+                if env_value is not None:
+                    env_bound[name] = env_value
+        if env_bound:
+            symbols["Env"] = env_bound
         # Debug log the Local symbols to help diagnose type issues
         if local_data:
             for key, value in local_data.items():
@@ -877,6 +972,11 @@ class DeclarativeActionExecutor(Executor):
         action_id = id or action_def.get("id") or f"{action_def.get('kind', 'action')}_{hash(str(action_def)) % 10000}"
         super().__init__(id=action_id, defer_discovery=True)
         self._action_def = action_def
+        # The active :class:`DeclarativeEnvConfig` is stamped onto the
+        # executor by :class:`DeclarativeWorkflowBuilder` after construction.
+        # Defaults to an empty configuration so direct ``DeclarativeActionExecutor``
+        # construction (e.g. in unit tests) doesn't expose ``os.environ``.
+        self._declarative_env_config: DeclarativeEnvConfig = _EMPTY_ENV_CONFIG
 
         # Manually register handlers after initialization
         self._handlers = {}
@@ -884,6 +984,16 @@ class DeclarativeActionExecutor(Executor):
         self._discover_handlers()
         self._discover_response_handlers()
 
+    def set_declarative_env_config(self, env_config: DeclarativeEnvConfig) -> None:
+        """Set the workflow-level :class:`DeclarativeEnvConfig` for this executor.
+
+        Called by :class:`DeclarativeWorkflowBuilder` after each executor is
+        created so that ``_to_powerfx_symbols`` populates the ``Env`` symbol
+        according to the caller-supplied configuration on the
+        :class:`WorkflowFactory`.
+        """
+        self._declarative_env_config = env_config
+
     @property
     def action_def(self) -> dict[str, Any]:
         """Get the action definition."""
@@ -896,7 +1006,7 @@ class DeclarativeActionExecutor(Executor):
 
     def _get_state(self, state: State) -> DeclarativeWorkflowState:
         """Get the declarative workflow state wrapper."""
-        return DeclarativeWorkflowState(state)
+        return DeclarativeWorkflowState(state, env_config=self._declarative_env_config)
 
     async def _ensure_state_initialized(
         self,

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

@@ -22,8 +22,10 @@ from agent_framework import (
 )
 
 from ._declarative_base import (
+    _EMPTY_ENV_CONFIG,  # type: ignore[reportPrivateUsage]
     ConditionResult,
     DeclarativeActionExecutor,
+    DeclarativeEnvConfig,
     LoopIterationResult,
 )
 from ._errors import DeclarativeWorkflowError
@@ -140,6 +142,7 @@ class DeclarativeWorkflowBuilder:
         max_iterations: int | None = None,
         http_request_handler: HttpRequestHandler | None = None,
         mcp_tool_handler: MCPToolHandler | None = None,
+        env_config: DeclarativeEnvConfig | None = None,
     ):
         """Initialize the builder.
 
@@ -158,6 +161,10 @@ class DeclarativeWorkflowBuilder:
             mcp_tool_handler: Handler used to dispatch InvokeMcpTool calls.
                 Must be supplied when the workflow contains any InvokeMcpTool;
                 otherwise build raises ``DeclarativeWorkflowError``.
+            env_config: Optional :class:`DeclarativeEnvConfig` controlling
+                how the ``Env`` PowerFx symbol is populated for every
+                executor built by this builder. Defaults to an empty
+                configuration (``Env`` not exposed).
         """
         self._yaml_def = yaml_definition
         self._workflow_id = workflow_id or yaml_definition.get("name", "declarative_workflow")
@@ -171,6 +178,7 @@ class DeclarativeWorkflowBuilder:
         self._seen_explicit_ids: set[str] = set()  # Track explicit IDs for duplicate detection
         self._http_request_handler = http_request_handler
         self._mcp_tool_handler = mcp_tool_handler
+        self._env_config: DeclarativeEnvConfig = env_config if env_config is not None else _EMPTY_ENV_CONFIG
         # Resolve max_iterations: explicit arg > YAML maxTurns > core default
         resolved = max_iterations if max_iterations is not None else yaml_definition.get("maxTurns")
         if resolved is not None and (not isinstance(resolved, int) or resolved <= 0):
@@ -221,6 +229,15 @@ class DeclarativeWorkflowBuilder:
         # Resolve pending gotos (back-edges for loops, forward-edges for jumps)
         self._resolve_pending_gotos(builder)
 
+        # Stamp the resolved DeclarativeEnvConfig onto every executor so they
+        # expose the configured Env binding through their _get_state(). This
+        # happens after _create_executors_for_actions and _resolve_pending_gotos
+        # so it covers the entry node, join nodes, evaluators, foreach
+        # init/next/exit nodes, and goto placeholders.
+        for executor in self._executors.values():
+            if isinstance(executor, DeclarativeActionExecutor):
+                executor.set_declarative_env_config(self._env_config)
+
         return builder.build()
 
     def _validate_workflow(self, actions: list[dict[str, Any]]) -> None:

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

@@ -26,7 +26,7 @@ from agent_framework import (
 )
 
 from .._loader import AgentFactory
-from .._models import _safe_mode_context  # type: ignore[reportPrivateUsage]
+from ._declarative_base import DeclarativeEnvConfig, discover_env_references
 from ._declarative_builder import DeclarativeWorkflowBuilder
 from ._errors import DeclarativeWorkflowError
 from ._http_handler import HttpRequestHandler
@@ -94,7 +94,8 @@ class WorkflowFactory:
         max_iterations: int | None = None,
         http_request_handler: HttpRequestHandler | None = None,
         mcp_tool_handler: MCPToolHandler | None = None,
-        safe_mode: bool = True,
+        configuration: Mapping[str, str] | None = None,
+        restrict_env_to_configuration: bool = True,
     ) -> None:
         """Initialize the workflow factory.
 
@@ -121,15 +122,25 @@ 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.
+            configuration: Optional mapping that populates the PowerFx ``Env``
+                symbol referenced from workflow YAML expressions (e.g.
+                ``=Env.MY_KEY``). Mirrors the .NET ``IConfiguration``-based
+                pattern in
+                ``dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowDiagnostics.cs``.
+                Keys supplied here are always exposed under ``Env.<key>``; the
+                process ``os.environ`` is consulted only when
+                ``restrict_env_to_configuration`` is ``False``. When neither
+                source produces a value the ``Env`` symbol is omitted so
+                ``=Env.X`` evaluates to the literal expression string.
+            restrict_env_to_configuration: When ``True`` (default), the
+                ``Env`` PowerFx symbol is populated exclusively from
+                ``configuration``; ``os.environ`` is never consulted. Set to
+                ``False`` to additionally fall back to ``os.environ`` for
+                names absent from ``configuration`` that the workflow YAML
+                explicitly references. The fallback is constrained to names
+                discovered in PowerFx expressions inside the workflow
+                definition so unrelated environment variables never enter
+                the PowerFx scope.
 
         Examples:
             .. code-block:: python
@@ -162,9 +173,20 @@ class WorkflowFactory:
                     checkpoint_storage=FileCheckpointStorage("./checkpoints"),
                     env_file=".env",
                 )
+
+            .. code-block:: python
+
+                from agent_framework.declarative import WorkflowFactory
+
+                # Inject named values for =Env.* references in the workflow YAML
+                factory = WorkflowFactory(
+                    configuration={
+                        "MY_SERVER_URL": "https://example.com",
+                        "MY_TOOL_NAME": "search",
+                    },
+                )
         """
-        self.safe_mode = safe_mode
-        self._agent_factory = agent_factory or AgentFactory(env_file_path=env_file, safe_mode=safe_mode)
+        self._agent_factory = agent_factory or AgentFactory(env_file_path=env_file)
         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
@@ -172,6 +194,8 @@ class WorkflowFactory:
         self._max_iterations = max_iterations
         self._http_request_handler = http_request_handler
         self._mcp_tool_handler = mcp_tool_handler
+        self._configuration: dict[str, str] = dict(configuration) if configuration else {}
+        self._restrict_env_to_configuration = restrict_env_to_configuration
 
     def create_workflow_from_yaml_path(
         self,
@@ -350,15 +374,6 @@ 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", "")
@@ -415,6 +430,16 @@ class WorkflowFactory:
         if description:
             normalized_def["description"] = description
 
+        # Build the DeclarativeEnvConfig from the factory's configuration and the
+        # set of Env references actually used in the workflow PowerFx expressions.
+        # The referenced-name allowlist constrains ``os.environ`` fallback (when
+        # enabled) so unrelated variables never enter the PowerFx scope.
+        env_config = DeclarativeEnvConfig(
+            values=dict(self._configuration),
+            restrict_to_configuration=self._restrict_env_to_configuration,
+            referenced_names=frozenset(discover_env_references(normalized_def)),
+        )
+
         # Build the graph-based workflow, passing agents and tools for specialized executors
         try:
             graph_builder = DeclarativeWorkflowBuilder(
@@ -426,6 +451,7 @@ class WorkflowFactory:
                 max_iterations=self._max_iterations,
                 http_request_handler=self._http_request_handler,
                 mcp_tool_handler=self._mcp_tool_handler,
+                env_config=env_config,
             )
             workflow = graph_builder.build()
         except ValueError as e:

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

@@ -33,7 +33,7 @@ import logging
 from collections import OrderedDict
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Protocol, cast, runtime_checkable
+from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast, runtime_checkable
 
 import httpx
 
@@ -194,6 +194,20 @@ class DefaultMCPToolHandler:
             Defaults to ``32``.
     """
 
+    LIST_TOOLS_TOOL_NAME: ClassVar[str] = "tools/list"
+    """Reserved ``tool_name`` that maps an :class:`MCPToolHandler` invocation
+    to the MCP protocol ``tools/list`` discovery operation.
+
+    Mirrors the .NET ``DefaultMcpToolHandler.ListToolsToolName`` public
+    constant for cross-language discoverability. When this handler receives
+    an invocation with this name it pages through ``session.list_tools()``
+    and returns the catalog as a single ``TextContent`` containing JSON of
+    shape ``{"tools": [{name, description, inputSchema, outputSchema}, ...]}``.
+    Workflows can reference this name from an ``InvokeMcpTool`` declarative
+    action to introspect a server's tool surface without an extra round-trip
+    from host code.
+    """
+
     def __init__(
         self,
         *,
@@ -217,10 +231,29 @@ class DefaultMCPToolHandler:
         self._closed = False
 
     async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult:
-        """Invoke ``invocation.tool_name`` on the cached MCP client for the server."""
+        """Invoke ``invocation.tool_name`` on the cached MCP client for the server.
+
+        The reserved name :attr:`LIST_TOOLS_TOOL_NAME` (``"tools/list"``) is
+        intercepted client-side: instead of being forwarded as a tool call,
+        it is translated to an MCP ``session.list_tools()`` discovery
+        operation (paginated automatically) and returned as a single
+        ``TextContent`` containing a JSON tool catalog. Matches the .NET
+        ``DefaultMcpToolHandler`` behaviour so the same YAML works
+        cross-language.
+        """
         from agent_framework import Content
         from agent_framework.exceptions import ToolExecutionException
 
+        # Reserved-name args validation runs before connect: rejecting bad
+        # input shouldn't require establishing an MCP session.
+        if invocation.tool_name == self.LIST_TOOLS_TOOL_NAME and invocation.arguments:
+            message = f"The reserved MCP '{self.LIST_TOOLS_TOOL_NAME}' operation does not accept tool arguments."
+            return MCPToolResult(
+                outputs=[Content.from_text(f"Error: {message}")],
+                is_error=True,
+                error_message=message,
+            )
+
         try:
             entry = await self._get_or_create_entry(invocation)
         except Exception as exc:
@@ -240,6 +273,8 @@ class DefaultMCPToolHandler:
             )
 
         try:
+            if invocation.tool_name == self.LIST_TOOLS_TOOL_NAME:
+                return await self._invoke_list_tools(entry)
             raw = await entry.tool.call_tool(invocation.tool_name, **invocation.arguments)
         except ToolExecutionException as exc:
             logger.info(
@@ -284,6 +319,60 @@ class DefaultMCPToolHandler:
             outputs = list(raw)
         return MCPToolResult(outputs=outputs)
 
+    @staticmethod
+    async def _invoke_list_tools(entry: _CacheEntry) -> MCPToolResult:
+        """Handle the reserved :attr:`LIST_TOOLS_TOOL_NAME` invocation.
+
+        Pages through ``session.list_tools()`` (mirroring the pagination loop
+        in :meth:`agent_framework.MCPTool.load_tools`) and serialises the
+        full catalog as a single ``TextContent`` containing JSON of shape
+        ``{"tools": [{name, description, inputSchema, outputSchema}, ...]}``.
+
+        The output shape, property names, and property order match
+        ``DefaultMcpToolHandler.SerializeToolsList`` in the .NET reference
+        implementation so the same workflow YAML can consume the result
+        cross-language. ``indent=2`` matches ``Utf8JsonWriter`` ``Indented``
+        mode; ``allow_nan=False`` guards against producing non-conformant
+        JSON ``NaN``/``Infinity`` tokens if a misbehaving server returns
+        such values in a schema.
+        """
+        from agent_framework import Content
+
+        session = getattr(entry.tool, "session", None)
+        if session is None:
+            message = "MCP session is not connected; cannot list tools."
+            return MCPToolResult(
+                outputs=[Content.from_text(f"Error: {message}")],
+                is_error=True,
+                error_message=message,
+            )
+
+        # Lazy import keeps ``mcp`` types out of module import time.
+        from mcp import types as mcp_types
+
+        collected: list[Any] = []
+        params: mcp_types.PaginatedRequestParams | None = None
+        while True:
+            tool_list = await session.list_tools(params=params)
+            collected.extend(tool_list.tools)
+            next_cursor = getattr(tool_list, "nextCursor", None)
+            if not next_cursor:
+                break
+            params = mcp_types.PaginatedRequestParams(cursor=next_cursor)
+
+        payload = {
+            "tools": [
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "inputSchema": tool.inputSchema,
+                    "outputSchema": tool.outputSchema,
+                }
+                for tool in collected
+            ],
+        }
+        return MCPToolResult(outputs=[Content.from_text(json.dumps(payload, indent=2, allow_nan=False))])
+
     async def aclose(self) -> None:
         """Close all cached MCP clients and the owned httpx clients.
 

python/packages/declarative/tests/test_default_mcp_tool_handler.py

@@ -13,6 +13,7 @@ owned-vs-caller httpx close semantics.
 from __future__ import annotations
 
 import asyncio
+import json
 import sys
 from typing import Any
 from unittest.mock import patch
@@ -33,6 +34,55 @@ pytestmark = pytest.mark.skipif(
 )
 
 
+class FakeListToolsResult:  # noqa: B903 - mimics ``mcp.types.ListToolsResult`` shape, not a value type
+    """Stand-in for ``mcp.types.ListToolsResult`` returned by ``session.list_tools()``."""
+
+    def __init__(self, tools: list[Any], next_cursor: str | None = None) -> None:
+        self.tools = tools
+        self.nextCursor = next_cursor
+
+
+class FakeMcpTool:
+    """Stand-in for an MCP ``Tool`` (subset used by ``_invoke_list_tools``)."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str | None = None,
+        inputSchema: dict[str, Any] | None = None,
+        outputSchema: dict[str, Any] | None = None,
+    ) -> None:
+        self.name = name
+        self.description = description
+        self.inputSchema = inputSchema if inputSchema is not None else {"type": "object", "properties": {}}
+        self.outputSchema = outputSchema
+
+
+class FakeMcpSession:
+    """Stand-in for ``mcp.ClientSession``.
+
+    ``list_tools_pages`` lets a test enqueue multiple paginated responses;
+    when None (default), an empty single-page result is returned. ``list_tools_error``
+    raises a synthetic error on the next call when set.
+    """
+
+    def __init__(self) -> None:
+        self.list_tools_pages: list[FakeListToolsResult] | None = None
+        self.list_tools_calls: list[Any] = []
+        self.list_tools_error: BaseException | None = None
+
+    async def list_tools(self, params: Any = None) -> FakeListToolsResult:
+        self.list_tools_calls.append(params)
+        if self.list_tools_error is not None:
+            raise self.list_tools_error
+        if self.list_tools_pages is None:
+            return FakeListToolsResult(tools=[])
+        index = len(self.list_tools_calls) - 1
+        if index >= len(self.list_tools_pages):
+            return FakeListToolsResult(tools=[])
+        return self.list_tools_pages[index]
+
+
 class FakeTool:
     """Stand-in for ``MCPStreamableHTTPTool``.
 
@@ -50,6 +100,7 @@ class FakeTool:
         self.connect_error: BaseException | None = None
         self.call_handler: Any = lambda **_a: [Content.from_text("ok")]
         self._httpx_client: httpx.AsyncClient | None = None
+        self.session: FakeMcpSession | None = None
         # Mimic MCPStreamableHTTPTool: when no caller client AND header_provider
         # is set, lazily allocate an owned httpx client during connect.
         FakeTool.instances.append(self)
@@ -63,6 +114,9 @@ class FakeTool:
         # Mimic lazy httpx allocation when no client provided AND header_provider set.
         if self.kwargs.get("http_client") is None and self.kwargs.get("header_provider") is not None:
             self._httpx_client = httpx.AsyncClient()
+        # Mimic MCPStreamableHTTPTool: a live session becomes available after connect.
+        if self.session is None:
+            self.session = FakeMcpSession()
 
     async def close(self) -> None:
         self.close_count += 1
@@ -541,3 +595,185 @@ class TestCacheKey:
         k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"X": "Bearer-A"})
         k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"X": "bearer-a"})
         assert k1 != k2
+
+
+# ---------- tools/list reserved name --------------------------------------
+
+
+class TestListTools:
+    """Exercise the reserved :attr:`DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME` interception path."""
+
+    @pytest.mark.asyncio
+    async def test_list_tools_returns_json_catalog(self) -> None:
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            # Prime the cache so the FakeTool session exists.
+            await handler.invoke_tool(_invocation())
+            FakeTool.instances[0].session.list_tools_pages = [  # type: ignore[union-attr]
+                FakeListToolsResult(
+                    tools=[
+                        FakeMcpTool(
+                            name="search",
+                            description="Search docs",
+                            inputSchema={"type": "object", "properties": {"q": {"type": "string"}}},
+                            outputSchema={"type": "object"},
+                        ),
+                        FakeMcpTool(name="echo", description=None, outputSchema=None),
+                    ],
+                ),
+            ]
+            result = await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        assert result.is_error is False
+        assert len(result.outputs) == 1
+        payload = json.loads(result.outputs[0].text)  # type: ignore[reportAttributeAccessIssue]
+        assert payload == {
+            "tools": [
+                {
+                    "name": "search",
+                    "description": "Search docs",
+                    "inputSchema": {"type": "object", "properties": {"q": {"type": "string"}}},
+                    "outputSchema": {"type": "object"},
+                },
+                {
+                    "name": "echo",
+                    "description": None,
+                    "inputSchema": {"type": "object", "properties": {}},
+                    "outputSchema": None,
+                },
+            ],
+        }
+
+    @pytest.mark.asyncio
+    async def test_list_tools_property_order_matches_dotnet(self) -> None:
+        """The JSON property order must match .NET SerializeToolsList: name, description, inputSchema, outputSchema."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            await handler.invoke_tool(_invocation())
+            FakeTool.instances[0].session.list_tools_pages = [  # type: ignore[union-attr]
+                FakeListToolsResult(tools=[FakeMcpTool(name="t1", description="d")]),
+            ]
+            result = await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        text = result.outputs[0].text  # type: ignore[reportAttributeAccessIssue]
+        name_idx = text.find('"name"')
+        desc_idx = text.find('"description"')
+        input_idx = text.find('"inputSchema"')
+        output_idx = text.find('"outputSchema"')
+        assert 0 <= name_idx < desc_idx < input_idx < output_idx
+
+    @pytest.mark.asyncio
+    async def test_list_tools_indented_output(self) -> None:
+        """Output is indented with 2-space indent (matches .NET ``Indented=true``)."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            await handler.invoke_tool(_invocation())
+            FakeTool.instances[0].session.list_tools_pages = [  # type: ignore[union-attr]
+                FakeListToolsResult(tools=[FakeMcpTool(name="t1")]),
+            ]
+            result = await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        text = result.outputs[0].text  # type: ignore[reportAttributeAccessIssue]
+        # Indented output contains newlines and a 2-space indented key.
+        assert "\n  " in text
+
+    @pytest.mark.asyncio
+    async def test_list_tools_rejects_arguments(self) -> None:
+        """Reserved name does NOT accept tool arguments. Fails fast before connect."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            result = await handler.invoke_tool(
+                _invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME, arguments={"q": "test"}),
+            )
+        assert result.is_error is True
+        assert "does not accept tool arguments" in (result.error_message or "")
+        # Args validation runs before connect, so no tool was instantiated.
+        assert FakeTool.instances == []
+
+    @pytest.mark.asyncio
+    async def test_list_tools_empty_args_dict_is_accepted(self) -> None:
+        """An empty arguments dict is equivalent to no arguments."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            await handler.invoke_tool(_invocation())
+            result = await handler.invoke_tool(
+                _invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME, arguments={}),
+            )
+        assert result.is_error is False
+
+    @pytest.mark.asyncio
+    async def test_list_tools_paginates(self) -> None:
+        """Pagination loop calls list_tools repeatedly until nextCursor is empty."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            await handler.invoke_tool(_invocation())
+            FakeTool.instances[0].session.list_tools_pages = [  # type: ignore[union-attr]
+                FakeListToolsResult(tools=[FakeMcpTool(name="a")], next_cursor="cursor1"),
+                FakeListToolsResult(tools=[FakeMcpTool(name="b")], next_cursor="cursor2"),
+                FakeListToolsResult(tools=[FakeMcpTool(name="c")], next_cursor=None),
+            ]
+            result = await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        payload = json.loads(result.outputs[0].text)  # type: ignore[reportAttributeAccessIssue]
+        assert [t["name"] for t in payload["tools"]] == ["a", "b", "c"]
+        session = FakeTool.instances[0].session
+        assert session is not None
+        assert len(session.list_tools_calls) == 3
+        # First call has no cursor; second/third use the cursor from the prior page.
+        assert session.list_tools_calls[0] is None
+        assert getattr(session.list_tools_calls[1], "cursor", None) == "cursor1"
+        assert getattr(session.list_tools_calls[2], "cursor", None) == "cursor2"
+
+    @pytest.mark.asyncio
+    async def test_list_tools_shares_cache_with_call_tool(self) -> None:
+        """tools/list reuses the same cached MCP session as a regular call_tool."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            await handler.invoke_tool(_invocation(tool_name="search"))
+            await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        assert len(FakeTool.instances) == 1
+        assert FakeTool.instances[0].connect_count == 1
+
+    @pytest.mark.asyncio
+    async def test_list_tools_propagates_session_errors_as_error_result(self) -> None:
+        """Errors raised by session.list_tools become MCPToolResult(is_error=True), not crashes."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            await handler.invoke_tool(_invocation())
+            FakeTool.instances[0].session.list_tools_error = httpx.ReadTimeout("read timed out")  # type: ignore[union-attr]
+            result = await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        assert result.is_error is True
+        assert "ReadTimeout" in (result.error_message or "")
+
+    @pytest.mark.asyncio
+    async def test_list_tools_returns_error_when_session_is_none(self) -> None:
+        """If somehow the cached tool has no session, return a clear error rather than crashing."""
+        handler = DefaultMCPToolHandler()
+        with _patch_tool():
+            await handler.invoke_tool(_invocation())
+            FakeTool.instances[0].session = None
+            result = await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        assert result.is_error is True
+        assert "not connected" in (result.error_message or "")
+
+    @pytest.mark.asyncio
+    async def test_list_tools_does_not_call_call_tool(self) -> None:
+        """The reserved name is intercepted; the inner call_tool path is bypassed."""
+        handler = DefaultMCPToolHandler()
+        call_tool_invoked = False
+
+        def fail(**_a: Any) -> Any:
+            nonlocal call_tool_invoked
+            call_tool_invoked = True
+            raise AssertionError("call_tool should not run for tools/list")
+
+        with _patch_tool():
+            await handler.invoke_tool(_invocation())
+            FakeTool.instances[0].call_handler = fail
+            FakeTool.instances[0].session.list_tools_pages = [  # type: ignore[union-attr]
+                FakeListToolsResult(tools=[]),
+            ]
+            result = await handler.invoke_tool(_invocation(tool_name=DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME))
+        assert call_tool_invoked is False
+        assert result.is_error is False
+
+    def test_class_attribute_value(self) -> None:
+        # Constant name MUST be the MCP protocol method name for cross-language parity
+        # with .NET ``DefaultMcpToolHandler.ListToolsToolName``.
+        assert DefaultMCPToolHandler.LIST_TOOLS_TOOL_NAME == "tools/list"

python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/main.py

@@ -11,35 +11,32 @@ tools (e.g. ``web_search``) returns the tool under its plain name.
 
 This sample mirrors the .NET sample
 ``dotnet/samples/03-workflows/Declarative/InvokeFoundryToolboxMcp/`` and
-shows how to:
+focuses on the **workflow execution path**:
 
-  1. Provision a toolbox in a Foundry project (delete-then-create_version,
-     so the sample can be re-run without manual cleanup).
-  2. Configure a ``WorkflowFactory`` with a custom :class:`MCPToolHandler`
-     that:
-       * routes every MCP request through a single
-         :class:`httpx.AsyncClient` carrying an Azure AD bearer token
-         (the toolbox endpoint requires AAD auth), and
-       * intercepts the reserved tool name ``"tools/list"`` so the YAML
-         can introspect the toolbox tool set without an extra Python
-         round-trip (matching the .NET ``DefaultMcpToolHandler``
-         behaviour).
-  3. Invoke ``microsoft_docs_search`` (from the Microsoft Learn Docs MCP
-     server surfaced by the toolbox) and ``web_search`` (Foundry built-in)
-     from a single declarative workflow.
-  4. Hand both result sets to a local :class:`Agent` registered with the
-     factory by name so the workflow's ``InvokeAzureAgent`` action can
-     summarise them.
+  1. Build a bearer-authenticated ``httpx.AsyncClient`` for the toolbox
+     MCP proxy and hand it to :class:`DefaultMCPToolHandler` so the YAML
+     can call MCP tools (and introspect the toolbox tool list via the
+     reserved ``"tools/list"`` tool name handled natively by the
+     framework, matching .NET
+     ``DefaultMcpToolHandler.ListToolsToolName``).
+  2. Configure a :class:`WorkflowFactory` with that handler plus a local
+     :class:`Agent` registered by name so the YAML's ``InvokeAzureAgent``
+     action can summarise the combined tool output.
+  3. Drive the workflow with a user question and render per-action
+     progress markers plus the final agent summary.
+
+One-off **toolbox administration** (delete + create_version) is delegated
+to :mod:`toolbox_provisioning` so this file stays focused on the workflow.
 
 Security note:
     The default ``DefaultMCPToolHandler`` performs no URL allowlisting or
-    SSRF protection. This sample wraps it with a project-scoped handler
+    SSRF protection. This sample uses a project-scoped ``client_provider``
     that pins outbound requests to ``Authorization: Bearer …`` via Azure
-    AD; for production deployments, additionally constrain the workflow
-    YAML to a known toolbox URL and reject any other server URL before
-    delegating to the inner handler. MCP outputs flow back into agent
-    conversations and share the prompt-injection risk surface of any
-    other tool output.
+    AD AND fails closed (raises) when the YAML resolves a different
+    ``serverUrl``, so a tampered ``=Env.*`` value cannot redirect the
+    bearer token to an attacker-controlled URL. MCP outputs flow back
+    into agent conversations and share the prompt-injection risk
+    surface of any other tool output.
 
 Run with:
     python samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/main.py
@@ -90,31 +87,31 @@ Sample output:
 """
 
 import asyncio
-import contextlib
-import json
 import os
-import sys
 from collections.abc import Iterator
 from pathlib import Path
-from typing import Any
 
 import httpx
-from agent_framework import Agent, Content, MCPStreamableHTTPTool
+from agent_framework import Agent
 from agent_framework.declarative import (
     DefaultMCPToolHandler,
     MCPToolInvocation,
-    MCPToolResult,
     WorkflowFactory,
 )
 from agent_framework.foundry import FoundryChatClient
 from azure.core.credentials import TokenCredential
 from azure.identity import AzureCliCredential, get_bearer_token_provider
+from toolbox_provisioning import (
+    AZ_CLI_PROCESS_TIMEOUT_SECONDS,
+    FOUNDRY_FEATURES_HEADERS,
+    build_toolbox_mcp_server_url,
+    create_sample_toolbox,
+)
 
 DEFAULT_TOOLBOX_NAME = "declarative_foundry_toolbox_mcp"
 DEFAULT_TOOLBOX_API_VERSION = "v1"
 DEFAULT_DOCS_SERVER_LABEL = "microsoft_docs"
 DEFAULT_WEB_SEARCH_TOOL_NAME = "web_search"
-DEFAULT_DOCS_MCP_SERVER_URL = "https://learn.microsoft.com/api/mcp"
 
 AGENT_NAME = "FoundryToolboxMcpAgent"
 
@@ -133,38 +130,10 @@ _ACTION_PROGRESS_LABELS: dict[str, str] = {
     SUMMARIZE_ACTION_ID: "Summarizing results...",
 }
 
-# Reserved tool name that the YAML uses to ask the handler for the toolbox
-# tool list. Mirrors .NET ``DefaultMcpToolHandler.ListToolsToolName``.
-LIST_TOOLS_TOOL_NAME = "tools/list"
-
 # AAD audience for the toolbox MCP proxy. Same scope used by the existing
 # Foundry hosted-toolbox samples.
 TOOLBOX_AAD_SCOPE = "https://ai.azure.com/.default"
 
-# Toolbox administration is gated by an Azure AI Foundry preview feature
-# flag. The .NET sample injects this header via a pipeline policy on the
-# ``AgentAdministrationClient``; the Python ``AIProjectClient`` doesn't
-# add it automatically, so we pass it as a per-call header on every
-# toolbox admin operation (delete + create_version) to make sure the
-# toolbox is actually provisioned in the V1Preview routing path that the
-# MCP proxy serves. Without this header, the calls can succeed at the
-# HTTP layer but the toolbox is never wired up to the MCP endpoint —
-# which surfaces at runtime as "MCP server failed to initialize:
-# Session terminated" on the first ``InvokeMcpTool`` call.
-FOUNDRY_FEATURES_HEADER_NAME = "Foundry-Features"
-FOUNDRY_FEATURES_HEADER_VALUE = "Toolboxes=V1Preview"
-FOUNDRY_FEATURES_HEADERS: dict[str, str] = {
-    FOUNDRY_FEATURES_HEADER_NAME: FOUNDRY_FEATURES_HEADER_VALUE,
-}
-
-# Bump the ``az.cmd`` subprocess timeout from the default 10s. On Windows
-# the Azure CLI batch wrapper can take noticeably longer than 10s to
-# return a token (cold-start + ``az`` self-update checks + AAD round-trip),
-# which surfaces as ``CredentialUnavailableError: Failed to invoke the
-# Azure CLI`` after a ``subprocess.TimeoutExpired`` from the credential's
-# internal call.
-AZ_CLI_PROCESS_TIMEOUT_SECONDS = 60
-
 # Match the MCP-recommended httpx timeouts (``mcp.shared._httpx_utils``:
 # 30s connect/write/pool, 5min SSE read). httpx's default ``Timeout(5.0)``
 # is far too aggressive for MCP streaming responses — long-running
@@ -189,70 +158,6 @@ result set contains an answer, say so plainly rather than guessing.
 """
 
 
-def build_toolbox_mcp_server_url(project_endpoint: str, name: str, api_version: str) -> str:
-    """Compose the Foundry toolbox MCP proxy URL.
-
-    Toolboxes provisioned via ``AIProjectClient.beta.toolboxes`` live under
-    the ``/toolboxes/{name}`` resource path (the Python SDK's
-    ``BetaToolboxesOperations`` routes POST/GET/DELETE there — see
-    ``azure/ai/projects/operations/_operations.py``). Their MCP proxy URL
-    is ``<project_endpoint>/toolboxes/{name}/mcp?api-version=<api_version>``,
-    matching the .NET sample.
-    """
-    base = project_endpoint.rstrip("/")
-    return f"{base}/toolboxes/{name}/mcp?api-version={api_version}"
-
-
-def create_sample_toolbox(
-    *,
-    name: str,
-    docs_server_label: str,
-    project_endpoint: str,
-    docs_server_url: str = DEFAULT_DOCS_MCP_SERVER_URL,
-) -> None:
-    """Provision a toolbox version in the Foundry project (idempotent).
-
-    Toolboxes are normally provisioned through the Foundry portal or a
-    deployment script; this helper exists so the sample can be re-run
-    end-to-end without manual cleanup. It deletes any toolbox under
-    ``name`` and then creates a new version that bundles:
-
-    - the Microsoft Learn Docs MCP server (``server_label=docs_server_label``),
-      and
-    - the Foundry built-in ``web_search`` tool.
-    """
-    from azure.ai.projects import AIProjectClient
-    from azure.ai.projects.models import MCPTool, Tool, WebSearchTool
-    from azure.core.exceptions import ResourceNotFoundError
-
-    with (
-        AzureCliCredential(process_timeout=AZ_CLI_PROCESS_TIMEOUT_SECONDS) as credential,
-        AIProjectClient(credential=credential, endpoint=project_endpoint) as project_client,
-    ):
-        try:
-            project_client.beta.toolboxes.delete(name, headers=FOUNDRY_FEATURES_HEADERS)
-            print(f"Toolbox '{name}' deleted (replacing with a fresh version).")
-        except ResourceNotFoundError:
-            pass
-
-        tools: list[Tool] = [
-            MCPTool(
-                server_label=docs_server_label,
-                server_url=docs_server_url,
-                require_approval="never",
-            ),
-            WebSearchTool(),
-        ]
-
-        created = project_client.beta.toolboxes.create_version(
-            name=name,
-            description="Sample toolbox combining Microsoft Learn Docs MCP and Foundry web search.",
-            tools=tools,
-            headers=FOUNDRY_FEATURES_HEADERS,
-        )
-        print(f"Created toolbox {created.name}@{created.version} ({len(created.tools)} tool(s)).")
-
-
 class _BearerAuth(httpx.Auth):
     """Inject a fresh Azure AD bearer token on every request.
 
@@ -269,94 +174,6 @@ class _BearerAuth(httpx.Auth):
         yield request
 
 
-class _ToolboxMcpToolHandler:
-    """:class:`MCPToolHandler` that adds ``tools/list`` support to the default handler.
-
-    The reserved tool name ``"tools/list"`` is intercepted client-side: it
-    is translated to an MCP ``session.list_tools()`` call and the result
-    is returned as a single JSON-encoded ``TextContent`` matching the
-    shape produced by the .NET ``DefaultMcpToolHandler``
-    (``{"tools": [{name, description, inputSchema, outputSchema}]}``).
-
-    All other tool invocations delegate to the wrapped
-    :class:`DefaultMCPToolHandler` so the LRU client cache, error
-    normalisation, and approval flow remain unchanged.
-
-    The ``tools/list`` path uses a transient :class:`MCPStreamableHTTPTool`
-    (``load_tools=False`` so MCP discovery only happens once via the
-    explicit ``session.list_tools()`` call). The same caller-supplied
-    ``httpx.AsyncClient`` is reused so the bearer token and any other
-    transport-level configuration stay consistent with the cached calls.
-    """
-
-    def __init__(self, inner: DefaultMCPToolHandler, http_client: httpx.AsyncClient) -> None:
-        self._inner = inner
-        self._http_client = http_client
-
-    async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult:
-        if invocation.tool_name == LIST_TOOLS_TOOL_NAME:
-            return await self._list_tools(invocation)
-        return await self._inner.invoke_tool(invocation)
-
-    async def _list_tools(self, invocation: MCPToolInvocation) -> MCPToolResult:
-        if invocation.arguments:
-            return MCPToolResult(
-                outputs=[Content.from_text("Error: 'tools/list' does not accept arguments.")],
-                is_error=True,
-                error_message="'tools/list' does not accept arguments.",
-            )
-
-        # Snapshot headers so the closure does not see later mutations.
-        captured_headers = dict(invocation.headers)
-
-        def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]:
-            return dict(captured_headers)
-
-        tool = MCPStreamableHTTPTool(
-            name=invocation.server_label or "foundry_toolbox_list",
-            url=invocation.server_url,
-            http_client=self._http_client,
-            header_provider=_header_provider if captured_headers else None,
-            load_tools=False,
-            load_prompts=False,
-        )
-
-        try:
-            await tool.connect()
-            tool_list = await tool.session.list_tools()  # type: ignore[union-attr]
-            payload = {
-                "tools": [
-                    {
-                        "name": entry.name,
-                        "description": entry.description,
-                        "inputSchema": entry.inputSchema,
-                        "outputSchema": entry.outputSchema,
-                    }
-                    for entry in tool_list.tools
-                ]
-            }
-            return MCPToolResult(outputs=[Content.from_text(json.dumps(payload))])
-        except Exception as exc:  # noqa: BLE001 - surface as tool error per protocol contract
-            message = f"{type(exc).__name__}: {exc}" if str(exc) else type(exc).__name__
-            return MCPToolResult(
-                outputs=[Content.from_text(f"Error: {message}")],
-                is_error=True,
-                error_message=message,
-            )
-        finally:
-            with contextlib.suppress(Exception):
-                await tool.close()
-
-    async def aclose(self) -> None:
-        await self._inner.aclose()
-
-    async def __aenter__(self) -> "_ToolboxMcpToolHandler":
-        return self
-
-    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
-        await self.aclose()
-
-
 async def main() -> None:
     """Run the Foundry toolbox MCP workflow."""
     # 1. Read configuration. ``FOUNDRY_PROJECT_ENDPOINT`` and
@@ -366,9 +183,7 @@ async def main() -> None:
     toolbox_name = os.environ.get("FOUNDRY_TOOLBOX_NAME", DEFAULT_TOOLBOX_NAME)
     toolbox_api_version = os.environ.get("FOUNDRY_TOOLBOX_API_VERSION", DEFAULT_TOOLBOX_API_VERSION)
     docs_server_label = os.environ.get("FOUNDRY_TOOLBOX_DOCS_SERVER_LABEL", DEFAULT_DOCS_SERVER_LABEL)
-    web_search_tool_name = os.environ.get(
-        "FOUNDRY_TOOLBOX_WEB_SEARCH_TOOL_NAME", DEFAULT_WEB_SEARCH_TOOL_NAME
-    )
+    web_search_tool_name = os.environ.get("FOUNDRY_TOOLBOX_WEB_SEARCH_TOOL_NAME", DEFAULT_WEB_SEARCH_TOOL_NAME)
 
     print("=" * 60)
     print("Invoke Foundry Toolbox MCP Workflow Demo")
@@ -382,17 +197,20 @@ async def main() -> None:
         project_endpoint=project_endpoint,
     )
 
-    # 3. Resolve the toolbox MCP proxy URL and publish all dynamic values
-    #    the YAML expects via ``Env.*``. Setting them after toolbox
-    #    creation ensures the URL points at the freshly created version.
+    # 3. Resolve the toolbox MCP proxy URL. The workflow YAML references
+    #    these values via ``=Env.FOUNDRY_TOOLBOX_*``; we publish them
+    #    through ``WorkflowFactory(configuration=...)`` so the values stay scoped to
+    #    this workflow.
     toolbox_endpoint = os.environ.get("FOUNDRY_TOOLBOX_ENDPOINT") or build_toolbox_mcp_server_url(
         project_endpoint=project_endpoint,
         name=toolbox_name,
         api_version=toolbox_api_version,
     )
-    os.environ["FOUNDRY_TOOLBOX_MCP_SERVER_URL"] = toolbox_endpoint
-    os.environ["FOUNDRY_TOOLBOX_DOCS_SERVER_LABEL"] = docs_server_label
-    os.environ["FOUNDRY_TOOLBOX_WEB_SEARCH_TOOL_NAME"] = web_search_tool_name
+    workflow_configuration: dict[str, str] = {
+        "FOUNDRY_TOOLBOX_MCP_SERVER_URL": toolbox_endpoint,
+        "FOUNDRY_TOOLBOX_DOCS_SERVER_LABEL": docs_server_label,
+        "FOUNDRY_TOOLBOX_WEB_SEARCH_TOOL_NAME": web_search_tool_name,
+    }
     print(f"Toolbox endpoint: {toolbox_endpoint}")
     print()
 
@@ -413,9 +231,9 @@ async def main() -> None:
 
     # 5. Build a bearer-authenticated httpx client. The same client is
     #    reused for every MCP request: the LRU cache inside
-    #    ``DefaultMCPToolHandler`` will keep a single MCP session alive
-    #    for the toolbox URL, and the ``tools/list`` interceptor reuses
-    #    the same httpx client so headers / auth stay consistent.
+    #    ``DefaultMCPToolHandler`` keeps a single MCP session alive
+    #    for the toolbox URL, and ``tools/list`` reuses that same
+    #    cached session for full transport-level consistency.
     #
     # Key configuration choices:
     #   * ``headers=FOUNDRY_FEATURES_HEADERS`` attaches the
@@ -451,36 +269,32 @@ async def main() -> None:
         # 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.
+        # ``Env.*`` value or a config injection), fail closed by raising so
+        # ``DefaultMCPToolHandler`` cannot fall back to an unauthenticated
+        # client that silently leaks the request shape.
         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,
+            raise ValueError(
+                f"Refusing to attach Foundry bearer token to unexpected MCP URL: "
+                f"{invocation.server_url!r}. Expected: {toolbox_endpoint!r}."
             )
-            return None
         return http_client
 
     async with (
         http_client,
-        DefaultMCPToolHandler(client_provider=_client_provider) as inner_handler,
-        _ToolboxMcpToolHandler(inner_handler, http_client) as mcp_handler,
+        DefaultMCPToolHandler(client_provider=_client_provider) as mcp_handler,
     ):
         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,
+            # the toolbox URL / tool names configurable without editing the
+            # YAML. We supply those values through ``configuration`` so the
+            # PowerFx ``Env`` symbol is populated from a local dict instead
+            # of the process environment. ``restrict_env_to_configuration``
+            # defaults to ``True`` which suppresses any ``os.environ``
+            # fallback — the workflow only sees the keys explicitly listed
+            # in ``workflow_configuration`` below.
+            configuration=workflow_configuration,
         )
 
         workflow_path = Path(__file__).parent / "workflow.yaml"

python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/toolbox_provisioning.py

@@ -0,0 +1,116 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Foundry toolbox provisioning helper for the ``invoke_foundry_toolbox_mcp`` sample.
+
+This module is intentionally narrow: it covers the one-off **administrative**
+setup needed to (re)create a Foundry toolbox so the sample can be run
+end-to-end without manual portal/CLI steps. Workflow execution, MCP
+handling, and agent orchestration live in :mod:`main`.
+
+Toolboxes are normally provisioned through the Foundry portal or a separate
+deployment script. Bundling the provisioning step here keeps the sample
+self-contained and re-runnable.
+
+The Foundry-Features preview header is exported here as well so the
+runtime MCP client in ``main.py`` can attach it on every outbound request
+(the MCP ``initialize`` handshake also requires the flag, not just the
+toolbox administration calls).
+"""
+
+from collections.abc import Mapping
+
+from azure.identity import AzureCliCredential
+
+DEFAULT_DOCS_MCP_SERVER_URL = "https://learn.microsoft.com/api/mcp"
+
+# Bump the ``az.cmd`` subprocess timeout from the default 10s. On Windows
+# the Azure CLI batch wrapper can take noticeably longer than 10s to
+# return a token (cold-start + ``az`` self-update checks + AAD round-trip),
+# which surfaces as ``CredentialUnavailableError: Failed to invoke the
+# Azure CLI`` after a ``subprocess.TimeoutExpired`` from the credential's
+# internal call.
+AZ_CLI_PROCESS_TIMEOUT_SECONDS = 60
+
+# Toolbox administration AND runtime MCP traffic are both gated by an
+# Azure AI Foundry preview feature flag. The .NET sample injects this
+# header via a pipeline policy on the ``AgentAdministrationClient``;
+# the Python ``AIProjectClient`` doesn't add it automatically, so we pass
+# it as a per-call header on every toolbox admin operation (delete +
+# create_version) here, and the runtime code in ``main.py`` attaches it
+# as a default header on the ``httpx.AsyncClient`` so it travels on the
+# MCP ``initialize`` handshake as well. Without this header on admin
+# calls, provisioning succeeds at the HTTP layer but the toolbox is
+# never wired up to the MCP endpoint — surfacing at runtime as "MCP
+# server failed to initialize: Session terminated" on the first
+# ``InvokeMcpTool`` call.
+FOUNDRY_FEATURES_HEADER_NAME = "Foundry-Features"
+FOUNDRY_FEATURES_HEADER_VALUE = "Toolboxes=V1Preview"
+FOUNDRY_FEATURES_HEADERS: Mapping[str, str] = {
+    FOUNDRY_FEATURES_HEADER_NAME: FOUNDRY_FEATURES_HEADER_VALUE,
+}
+
+
+def build_toolbox_mcp_server_url(project_endpoint: str, name: str, api_version: str) -> str:
+    """Compose the Foundry toolbox MCP proxy URL.
+
+    Toolboxes provisioned via ``AIProjectClient.beta.toolboxes`` live under
+    the ``/toolboxes/{name}`` resource path (the Python SDK's
+    ``BetaToolboxesOperations`` routes POST/GET/DELETE there — see
+    ``azure/ai/projects/operations/_operations.py``). Their MCP proxy URL
+    is ``<project_endpoint>/toolboxes/{name}/mcp?api-version=<api_version>``,
+    matching the .NET sample.
+    """
+    base = project_endpoint.rstrip("/")
+    return f"{base}/toolboxes/{name}/mcp?api-version={api_version}"
+
+
+def create_sample_toolbox(
+    *,
+    name: str,
+    docs_server_label: str,
+    project_endpoint: str,
+    docs_server_url: str = DEFAULT_DOCS_MCP_SERVER_URL,
+) -> None:
+    """Provision a toolbox version in the Foundry project (idempotent).
+
+    Deletes any existing toolbox under ``name`` and then creates a new
+    version that bundles:
+
+    - the Microsoft Learn Docs MCP server
+      (``server_label=docs_server_label``), and
+    - the Foundry built-in ``web_search`` tool.
+
+    Uses ``AzureCliCredential`` because the sample is meant to be run by
+    a developer with ``az login`` already configured; switch to a managed
+    identity / service principal credential for production deployments.
+    """
+    from azure.ai.projects import AIProjectClient
+    from azure.ai.projects.models import MCPTool, Tool, WebSearchTool
+    from azure.core.exceptions import ResourceNotFoundError
+
+    with (
+        AzureCliCredential(process_timeout=AZ_CLI_PROCESS_TIMEOUT_SECONDS) as credential,
+        AIProjectClient(credential=credential, endpoint=project_endpoint) as project_client,
+    ):
+        try:
+            project_client.beta.toolboxes.delete(name, headers=FOUNDRY_FEATURES_HEADERS)
+            print(f"Toolbox '{name}' deleted (replacing with a fresh version).")
+        except ResourceNotFoundError:
+            pass
+
+        tools: list[Tool] = [
+            MCPTool(
+                server_label=docs_server_label,
+                server_url=docs_server_url,
+                require_approval="never",
+            ),
+            WebSearchTool(),
+        ]
+
+        created = project_client.beta.toolboxes.create_version(
+            name=name,
+            description="Sample toolbox combining Microsoft Learn Docs MCP and Foundry web search.",
+            tools=tools,
+            headers=FOUNDRY_FEATURES_HEADERS,
+        )
+        print(f"Created toolbox {created.name}@{created.version} ({len(created.tools)} tool(s)).")

python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/workflow.yaml

@@ -7,10 +7,9 @@
 # The workflow:
 #   1. Accepts a documentation / web search query as input.
 #   2. Lists the tools exposed by the toolbox using the reserved
-#      toolName: ``tools/list``. The Python ``DefaultMCPToolHandler``
-#      does not intercept this name on its own; the sample's host code
-#      wraps it with a small handler that translates ``tools/list`` to
-#      an MCP ``session.list_tools()`` call.
+#      toolName: ``tools/list``. ``DefaultMCPToolHandler`` intercepts
+#      this reserved name natively and translates it
+#      to an MCP ``session.list_tools()`` call, returning a JSON catalog.
 #   3. Invokes the Microsoft Learn ``microsoft_docs_search`` MCP tool
 #      surfaced by the toolbox. Tool names from MCP-server-backed
 #      toolbox tools are namespaced as ``<server_label>___<tool_name>``.
@@ -41,9 +40,7 @@ trigger:
       variable: Local.SearchQuery
       value: =Workflow.Inputs.text
 
-    # List the tools exposed by the toolbox MCP proxy. The sample's
-    # custom MCPToolHandler intercepts the reserved ``tools/list`` name
-    # and returns the toolbox tool list as JSON.
+    # List the tools exposed by the toolbox MCP proxy.
     #
     # We intentionally OMIT ``conversationId`` here: the tool list is
     # metadata for the demo, not useful context for the downstream