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 dd326f62f3..450676b456 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py +++ b/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.`` (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, diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py index 67b4a58273..400ee48896 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py +++ b/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: diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index e3866933f3..f43be8094d 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/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.``; 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: diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 658ce42c23..8daa60dc9c 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/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. diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 3a5c67e1d6..027158e096 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/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" diff --git a/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/__init__.py b/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/__init__.py new file mode 100644 index 0000000000..2a50eae894 --- /dev/null +++ b/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Microsoft. All rights reserved. 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 ceffb08393..da92c8ba06 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 @@ -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 ``/toolboxes/{name}/mcp?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" diff --git a/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/toolbox_provisioning.py b/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/toolbox_provisioning.py new file mode 100644 index 0000000000..7572dee454 --- /dev/null +++ b/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 ``/toolboxes/{name}/mcp?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)).") diff --git a/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/workflow.yaml b/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/workflow.yaml index d84b915e05..510146822b 100644 --- a/python/samples/03-workflows/declarative/invoke_foundry_toolbox_mcp/workflow.yaml +++ b/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 ``___``. @@ -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