diff --git a/python/.github/skills/python-development/SKILL.md b/python/.github/skills/python-development/SKILL.md index 4214c08bdb..cbec0c523f 100644 --- a/python/.github/skills/python-development/SKILL.md +++ b/python/.github/skills/python-development/SKILL.md @@ -25,6 +25,10 @@ Every `.py` file must start with: - Use `Mapping` instead of `MutableMapping` for read-only input parameters - Prefer `# type: ignore[...]` over unnecessary casts, or `isinstance` checks, when these are internally called and executed methods But make sure the ignore is specific for both mypy and pyright so that we don't miss other mistakes +- Internal private helpers may be used across `agent_framework*` modules when intentional; use a targeted + `# pyright: ignore[reportPrivateUsage]` instead of making the helper public just to satisfy pyright. +- Do not add trivial pass-through or one-line helper functions solely to appease typing. Prefer targeted ignores, + casts, or clearer annotations over adding runtime overhead without a design benefit. ## Function Parameters diff --git a/python/CODING_STANDARD.md b/python/CODING_STANDARD.md index 9762f89242..04e576016f 100644 --- a/python/CODING_STANDARD.md +++ b/python/CODING_STANDARD.md @@ -105,7 +105,12 @@ Use typing as a helper first and suppressions as a last resort: (`# pyright: ignore[reportGeneralTypeIssues]`), file-level is allowed if there is a compelling reason for it, that should be documented right beneath the ignore. Never change the global suppression flags unless the dev team okays it. - **Private usage boundary**: Accessing private members across `agent_framework*` packages can be acceptable for this - codebase, but private member usage for non-Agent Framework dependencies should remain flagged. + codebase, but private member usage for non-Agent Framework dependencies should remain flagged. Do not make an + internal helper public merely to satisfy pyright private-usage checks inside the package; use a targeted + `# pyright: ignore[reportPrivateUsage]` when the internal dependency is intentional. +- **Avoid typing-only wrappers**: Do not introduce trivial pass-through functions solely to satisfy typing or private + usage checks. Prefer a targeted ignore, cast, or clearer annotation over an extra one-line function that adds runtime + overhead without improving the design. ## Function Parameter Guidelines diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index e3bb353bce..496b9b6601 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -85,6 +85,7 @@ agent_framework/ - **`MCPStdioTool`** / **`MCPStreamableHTTPTool`** / **`MCPWebsocketTool`** - Transport-specific subclasses. - **Argument allowlist (`_prepare_call_kwargs`)** - Before each `tools/call`, kwargs are filtered to an **allowlist** built from the tool's declared parameters (`inputSchema.properties`) plus any user-configured extras. Framework runtime kwargs injected through the function-invocation pipeline (e.g. `thread`, `conversation_id`, `chat_options`, `options`, `response_format`) are stripped by default rather than forwarded. A tool that declares no usable `properties` (including schemas with `additionalProperties: true`) forwards only the configured extras. The `_MCP_FRAMEWORK_DENYLIST` is a safety net for framework-named params a server *declares* in its schema (those are dropped); names explicitly opted in via `additional_tool_argument_names` always win. The reserved `_meta` key is never forwarded as an argument; trusted caller/runtime `_meta` is validated as MCP request metadata, model-supplied `_meta` is discarded in generated MCP functions, and metadata precedence is caller/runtime < OpenTelemetry < tools/list metadata. - **`allowed_tools`** (constructor arg on all `MCPTool` subclasses) - Restricts exposed MCP tools by raw remote MCP tool identity. Prefixed local names remain accepted only when the raw remote name already matches its normalized form; normalized/local aliases do not authorize a different raw remote name. If multiple raw remote tool names map to the same local function name, tool loading raises `ToolExecutionException` instead of first-one-wins shadowing. +- **Progressive MCP disclosure** (`use_progressive_disclosure`, `always_load`) - When enabled on any `MCPTool` subclass, the initial model-facing surface is loader tools (`list_mcp_tools` / `load_tool` / `unload_tool`, prefixed by `tool_name_prefix` when configured) plus allowed tools selected by `always_load` and tools loaded earlier on the same `MCPTool` instance. `list_mcp_tools` only reports tools that pass `allowed_tools`; filtered tools are not listed or loadable. Loader tool names are reserved in progressive mode: remote MCP tools whose local generated name collides with a loader name are omitted from the initial/listed surface, and explicit `load_tool` calls return a model-visible message pointing callers to `tool_name_prefix` or excluding the colliding tool. `load_tool` accepts one tool name or a list of tool names and uses `FunctionInvocationContext.add_tools(...)` so the selected generated MCP `FunctionTool`s become available on the next function-calling iteration while keeping existing approval mode, argument filtering, header-provider runtime kwargs, result parsing, OTel, and task behavior. `unload_tool` accepts one dynamically loaded tool name or a list of names and removes them from the live tool list and persisted progressive surface, but it does not remove tools configured in `always_load`. Invalid `always_load` entries are ignored like unmatched `allowed_tools` entries. - **`additional_tool_argument_names`** (constructor arg on all `MCPTool` subclasses) - Opt extra argument names back into the allowlist. Accepts a `Sequence[str]` (applied to every tool) or a `Mapping[str, Sequence[str]]` keyed by **remote tool name**, where the reserved key `"*"` denotes global extras. It is configured only in user code at construction; there is **no per-call/runtime override**, so a model-issued tool call cannot change which names pass through. To use a server that accepts `additionalProperties: true`, list the extra names here and then either (1) manually extend that tool's `inputSchema` (via the `.functions` list after connecting) so the model is prompted to supply them, or (2) supply the values yourself via `function_invocation_kwargs`. If a normal forwarded argument name is supplied by both the model and `function_invocation_kwargs`, the model-supplied value wins; `_meta` is the exception and only trusted runtime/caller metadata is used. - **Sampling guardrails** (`sampling_callback`) - Passing `client=` advertises `SamplingCapability` so the server can send `sampling/createMessage`. Because remote servers are untrusted (confused-deputy risk), the default `sampling_callback` is **deny-by-default** and applies, in order: a per-session rate limit (`sampling_max_requests`, default `_DEFAULT_SAMPLING_MAX_REQUESTS`), an approval gate (`sampling_approval_callback`), and a `maxTokens` cap (`sampling_max_tokens`, default `_DEFAULT_SAMPLING_MAX_TOKENS`). The approval callback (constructor arg on all subclasses; exported type alias `SamplingApprovalCallback`) receives the raw `CreateMessageRequestParams`, may be sync or async, and must return truthy to approve. When it is `None` (the default) every sampling request is denied; pass `lambda params: True` to restore legacy auto-approve as an explicit opt-in. Requests and denials are logged at WARNING (content is not logged). The per-session counter resets in `_reset_session_state`. - **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields: diff --git a/python/packages/core/agent_framework/_feature_stage.py b/python/packages/core/agent_framework/_feature_stage.py index 87392713cc..0bb44b9568 100644 --- a/python/packages/core/agent_framework/_feature_stage.py +++ b/python/packages/core/agent_framework/_feature_stage.py @@ -226,15 +226,16 @@ def _resolve_user_frame() -> tuple[str, int, str] | None: def _warn_on_feature_use( *, stage: FeatureStageName, - feature_id: str, + feature_id: str | Enum, object_name: str, category: type[Warning], ) -> None: - warning_key = (category, feature_id) + normalized_feature_id = _normalize_feature_id(feature_id) + warning_key = (category, normalized_feature_id) if warning_key in _WARNED_FEATURES: return - message = _build_stage_warning_message(stage=stage, feature_id=feature_id, object_name=object_name) + message = _build_stage_warning_message(stage=stage, feature_id=normalized_feature_id, object_name=object_name) user_frame = _resolve_user_frame() if user_frame is None: # Last-resort fallback: emit at the immediate caller of this helper. diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 83cfed92e7..5b3195545b 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -10,6 +10,7 @@ import logging import re import sys +import warnings from abc import abstractmethod from collections.abc import Callable, Collection, Coroutine, Mapping, Sequence from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore @@ -22,7 +23,12 @@ from opentelemetry import propagate from opentelemetry import trace as otel_trace -from ._feature_stage import ExperimentalFeature, experimental +from ._feature_stage import ( + ExperimentalFeature, + ExperimentalWarning, + _warn_on_feature_use, # pyright: ignore[reportPrivateUsage] + experimental, +) from ._tools import FunctionTool from ._types import ( ChatOptions, @@ -71,6 +77,9 @@ class MCPSpecificApproval(TypedDict, total=False): _MCP_REMOTE_NAME_KEY = "_mcp_remote_name" _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name" +_MCP_PROGRESSIVE_LIST_TOOL_NAME = "list_mcp_tools" +_MCP_PROGRESSIVE_LOAD_TOOL_NAME = "load_tool" +_MCP_PROGRESSIVE_UNLOAD_TOOL_NAME = "unload_tool" # Reserved key in an ``additional_tool_argument_names`` mapping that applies its # values to every tool on the server rather than a single named tool. _MCP_GLOBAL_EXTRA_ARGS_KEY = "*" @@ -409,6 +418,8 @@ def __init__( additional_properties: dict[str, Any] | None = None, task_options: MCPTaskOptions | None = None, additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, + use_progressive_disclosure: bool = False, + always_load: Collection[str] | None = None, ) -> None: """Initialize the MCP Tool base. @@ -469,7 +480,21 @@ def __init__( in addition to each tool's declared parameters. A ``Sequence[str]`` applies to every tool; a ``Mapping[str, Sequence[str]]`` is keyed by remote tool name with ``"*"`` as a global key. See the transport subclasses for full details. + use_progressive_disclosure: When ``True``, expose discovery and loader tools plus + ``always_load`` tools initially, and let the model load and unload other allowed MCP tools on demand. + always_load: MCP tool names to keep visible from the start when progressive disclosure + is enabled. Names use the same safe matching rules as ``allowed_tools``; unmatched + entries are ignored. """ + if use_progressive_disclosure and not load_tools: + raise ValueError("use_progressive_disclosure=True requires load_tools=True.") + if use_progressive_disclosure: + _warn_on_feature_use( + stage="experimental", + feature_id=ExperimentalFeature.PROGRESSIVE_TOOLS, + object_name="MCP progressive disclosure", + category=ExperimentalWarning, + ) self.name = name self.description = description or "" self.approval_mode = approval_mode @@ -498,6 +523,11 @@ def __init__( self.sampling_max_requests = sampling_max_requests self._sampling_request_count = 0 self._functions: list[FunctionTool] = [] + self.use_progressive_disclosure = use_progressive_disclosure + self.always_load = always_load + self._always_load_names = set(always_load or ()) + self._progressive_loader_functions: list[FunctionTool] | None = None + self._progressive_loaded_tool_names: set[str] = set() self._tool_call_meta_by_name: dict[str, dict[str, Any]] = {} self._tool_task_support_by_name: dict[str, str] = {} self._tool_param_names_by_name: dict[str, set[str]] = {} @@ -807,24 +837,273 @@ def _prepare_message_for_mcp( @property def functions(self) -> list[FunctionTool]: """Get the list of functions that are allowed.""" + if self.use_progressive_disclosure: + return self._progressive_functions() + return self._filtered_functions() + + def _filtered_functions(self) -> list[FunctionTool]: + """Return loaded MCP functions after applying ``allowed_tools``.""" if self.allowed_tools is None: return self._functions allowed_names = set(self.allowed_tools) filtered_functions: list[FunctionTool] = [] for func in self._functions: + if self._function_matches_names(func, allowed_names): + filtered_functions.append(func) + return filtered_functions + + def _function_matches_names(self, func: FunctionTool, names: set[str]) -> bool: + """Return whether a generated MCP function matches a configured name set.""" + if not names: + return False + additional_properties = func.additional_properties or {} + normalized_name = additional_properties.get(_MCP_NORMALIZED_NAME_KEY) + remote_name = additional_properties.get(_MCP_REMOTE_NAME_KEY) + if not isinstance(normalized_name, str) or not isinstance(remote_name, str): + return False + candidate_names = _mcp_config_candidate_names( + local_name=func.name, + normalized_name=normalized_name, + remote_name=remote_name, + ) + return any(name in names for name in candidate_names) + + def _progressive_functions(self) -> list[FunctionTool]: + """Return the initial model-facing function list for progressive disclosure.""" + initial_functions = list(self._progressive_loader_tools()) + initial_functions.extend( + func + for func in self._filtered_functions() + if ( + self._function_matches_names(func, self._always_load_names) + or func.name in self._progressive_loaded_tool_names + ) + and not self._function_collides_with_progressive_loader(func) + ) + return initial_functions + + def _progressive_loader_names(self) -> set[str]: + """Return the local names reserved for progressive disclosure loader tools.""" + return { + _build_prefixed_mcp_name(_MCP_PROGRESSIVE_LIST_TOOL_NAME, self.tool_name_prefix), + _build_prefixed_mcp_name(_MCP_PROGRESSIVE_LOAD_TOOL_NAME, self.tool_name_prefix), + _build_prefixed_mcp_name(_MCP_PROGRESSIVE_UNLOAD_TOOL_NAME, self.tool_name_prefix), + } + + def _function_collides_with_progressive_loader(self, func: FunctionTool) -> bool: + """Return whether an MCP function's local name collides with a loader tool name.""" + return func.name in self._progressive_loader_names() + + def _progressive_loader_tools(self) -> list[FunctionTool]: + """Create or return the generated progressive disclosure loader tools.""" + if self._progressive_loader_functions is not None: + return self._progressive_loader_functions + + list_tool_name = _build_prefixed_mcp_name(_MCP_PROGRESSIVE_LIST_TOOL_NAME, self.tool_name_prefix) + load_tool_name = _build_prefixed_mcp_name(_MCP_PROGRESSIVE_LOAD_TOOL_NAME, self.tool_name_prefix) + unload_tool_name = _build_prefixed_mcp_name(_MCP_PROGRESSIVE_UNLOAD_TOOL_NAME, self.tool_name_prefix) + self._progressive_loader_functions = [ + FunctionTool( + func=self._list_progressive_mcp_tools, + name=list_tool_name, + description="List the MCP tools that can be loaded from this server.", + approval_mode="never_require", + input_model={"type": "object", "properties": {}}, + ), + FunctionTool( + func=self._load_progressive_mcp_tool, + name=load_tool_name, + description="Load an MCP tool from this server so it can be called on the next iteration.", + approval_mode="never_require", + input_model={ + "type": "object", + "properties": { + "tool": { + "oneOf": [ + {"type": "string"}, + {"type": "array", "items": {"type": "string"}}, + ], + "description": "The MCP tool name, or MCP tool names, to load.", + }, + }, + "required": ["tool"], + }, + ), + FunctionTool( + func=self._unload_progressive_mcp_tool, + name=unload_tool_name, + description="Unload an MCP tool from this server when it is no longer relevant.", + approval_mode="never_require", + input_model={ + "type": "object", + "properties": { + "tool": { + "oneOf": [ + {"type": "string"}, + {"type": "array", "items": {"type": "string"}}, + ], + "description": "The MCP tool name, or MCP tool names, to unload.", + }, + }, + "required": ["tool"], + }, + ), + ] + return self._progressive_loader_functions + + def _resolve_progressive_function(self, tool_name: str) -> FunctionTool: + """Resolve a progressive disclosure tool name against allowed MCP functions.""" + matches = [func for func in self._filtered_functions() if self._function_matches_names(func, {tool_name})] + matches.extend(func for func in self._filtered_functions() if func.name == tool_name and func not in matches) + if not matches: + available = ( + ", ".join( + func.name + for func in self._filtered_functions() + if not self._function_collides_with_progressive_loader(func) + ) + or "none" + ) + raise ToolExecutionException(f"MCP tool '{tool_name}' is not available. Available tools: {available}.") + if len(matches) > 1: + raise ToolExecutionException(f"MCP tool name '{tool_name}' is ambiguous.") + return matches[0] + + @staticmethod + def _progressive_tool_names(tool: str | Sequence[str]) -> list[str]: + """Normalize a progressive loader request to a tool-name list.""" + if isinstance(tool, str): + return [tool] + if not isinstance(tool, Sequence): + raise ToolExecutionException("Progressive MCP tool request must be a string or a list of strings.") + tool_names = list(tool) + if not all(isinstance(tool_name, str) for tool_name in tool_names): + raise ToolExecutionException("Progressive MCP tool request must contain only strings.") + return tool_names + + def _is_progressive_function_loaded( + self, + func: FunctionTool, + ctx: FunctionInvocationContext | None, + ) -> bool: + """Return whether a progressive MCP function is already present in the live tool list.""" + if self._function_matches_names(func, self._always_load_names): + return True + if func.name in self._progressive_loaded_tool_names: + return True + if ctx is None or ctx.tools is None: + return False + return any(tool_item is func for tool_item in ctx.tools) + + def _list_progressive_mcp_tools(self, ctx: FunctionInvocationContext) -> list[dict[str, Any]]: + """List allowed MCP tools that can be loaded progressively.""" + tools: list[dict[str, Any]] = [] + for func in self._filtered_functions(): + if self._function_collides_with_progressive_loader(func): + continue additional_properties = func.additional_properties or {} - normalized_name = additional_properties.get(_MCP_NORMALIZED_NAME_KEY) remote_name = additional_properties.get(_MCP_REMOTE_NAME_KEY) - if not isinstance(normalized_name, str) or not isinstance(remote_name, str): + tools.append({ + "name": func.name, + "remote_name": remote_name if isinstance(remote_name, str) else func.name, + "description": func.description, + "parameters": func.parameters(), + "approval_mode": func.approval_mode, + "loaded": self._is_progressive_function_loaded(func, ctx), + "always_loaded": self._function_matches_names(func, self._always_load_names), + }) + return tools + + async def _load_progressive_mcp_tool(self, ctx: FunctionInvocationContext, tool: str | Sequence[str]) -> str: + """Load an allowed MCP tool into the live function-calling tool list.""" + if ctx.tools is None: + raise ToolExecutionException("load_tool can only be used inside an agent function-calling run.") + messages: list[str] = [] + functions_to_load: list[FunctionTool] = [] + function_names_to_load: set[str] = set() + for tool_name in self._progressive_tool_names(tool): + try: + func = self._resolve_progressive_function(tool_name) + except ToolExecutionException as ex: + messages.append(str(ex)) continue - candidate_names = _mcp_config_candidate_names( - local_name=func.name, - normalized_name=normalized_name, - remote_name=remote_name, - ) - if any(name in allowed_names for name in candidate_names): - filtered_functions.append(func) - return filtered_functions + if self._function_collides_with_progressive_loader(func): + loader_names = ", ".join(sorted(self._progressive_loader_names())) + messages.append( + f"MCP tool '{func.name}' conflicts with progressive disclosure loader tool name(s): " + f"{loader_names}. Set tool_name_prefix or exclude the colliding MCP tool." + ) + continue + if any(tool_item is func for tool_item in ctx.tools): + self._progressive_loaded_tool_names.add(func.name) + messages.append(f"MCP tool '{func.name}' is already available.") + continue + if func.name in function_names_to_load: + messages.append(f"MCP tool '{func.name}' is already queued to load.") + continue + functions_to_load.append(func) + function_names_to_load.add(func.name) + messages.append(f"Loaded MCP tool '{func.name}'. It is available on the next model iteration.") + if not messages: + return "No MCP tools requested." + if not functions_to_load: + return "\n".join(messages) + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ExperimentalWarning) + ctx.add_tools(functions_to_load) + except ValueError as ex: + raise ToolExecutionException(str(ex), inner_exception=ex) from ex + self._progressive_loaded_tool_names.update(func.name for func in functions_to_load) + return "\n".join(messages) + + async def _unload_progressive_mcp_tool(self, ctx: FunctionInvocationContext, tool: str | Sequence[str]) -> str: + """Unload a progressively loaded MCP tool from the live function-calling tool list.""" + if ctx.tools is None: + raise ToolExecutionException("unload_tool can only be used inside an agent function-calling run.") + messages: list[str] = [] + function_names_to_unload: set[str] = set() + for tool_name in self._progressive_tool_names(tool): + if tool_name in self._progressive_loader_names(): + messages.append(f"MCP loader tool '{tool_name}' cannot be unloaded.") + continue + try: + func = self._resolve_progressive_function(tool_name) + except ToolExecutionException as ex: + messages.append(str(ex)) + continue + if self._function_collides_with_progressive_loader(func): + loader_names = ", ".join(sorted(self._progressive_loader_names())) + messages.append( + f"MCP tool '{func.name}' conflicts with progressive disclosure loader tool name(s): " + f"{loader_names}. Set tool_name_prefix or exclude the colliding MCP tool." + ) + continue + if self._function_matches_names(func, self._always_load_names): + messages.append(f"MCP tool '{func.name}' is configured in always_load and cannot be unloaded.") + continue + if func.name in function_names_to_unload: + messages.append(f"MCP tool '{func.name}' is already queued to unload.") + continue + if func.name not in self._progressive_loaded_tool_names and not any( + tool_item is func for tool_item in ctx.tools + ): + messages.append(f"MCP tool '{func.name}' is not currently loaded.") + continue + function_names_to_unload.add(func.name) + messages.append(f"Unloaded MCP tool '{func.name}'. It will be removed on the next model iteration.") + if not messages: + return "No MCP tools requested." + if not function_names_to_unload: + return "\n".join(messages) + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ExperimentalWarning) + ctx.remove_tools(list(function_names_to_unload)) + except RuntimeError as ex: + raise ToolExecutionException(str(ex), inner_exception=ex) from ex + self._progressive_loaded_tool_names.difference_update(function_names_to_unload) + return "\n".join(messages) async def _ensure_lifecycle_owner(self) -> None: async with self._lifecycle_lock: @@ -2432,6 +2711,8 @@ def __init__( description: str | None = None, approval_mode: (Literal["always_require", "never_require"] | MCPSpecificApproval | None) = None, allowed_tools: Collection[str] | None = None, + use_progressive_disclosure: bool = False, + always_load: Collection[str] | None = None, args: list[str] | None = None, env: dict[str, str] | None = None, encoding: str | None = None, @@ -2487,6 +2768,11 @@ def __init__( disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is useful as a runtime guard or when you want to load tool metadata for inspection without exposing the tools for invocation. + use_progressive_disclosure: When ``True``, expose discovery and loader tools plus + ``always_load`` tools initially, and let the model load and unload other allowed MCP tools on demand. + always_load: MCP tool names to keep visible from the start when progressive disclosure + is enabled. Names use the same safe matching rules as ``allowed_tools``; unmatched + entries are ignored. additional_properties: Additional properties. args: The arguments to pass to the command. env: The environment variables to set for the command. @@ -2525,6 +2811,8 @@ def __init__( description=description, approval_mode=approval_mode, allowed_tools=allowed_tools, + use_progressive_disclosure=use_progressive_disclosure, + always_load=always_load, tool_name_prefix=tool_name_prefix, additional_properties=additional_properties, session=session, @@ -2612,6 +2900,8 @@ def __init__( description: str | None = None, approval_mode: (Literal["always_require", "never_require"] | MCPSpecificApproval | None) = None, allowed_tools: Collection[str] | None = None, + use_progressive_disclosure: bool = False, + always_load: Collection[str] | None = None, terminate_on_close: bool | None = None, client: SupportsChatGetResponse | None = None, sampling_approval_callback: SamplingApprovalCallback | None = None, @@ -2668,6 +2958,11 @@ def __init__( disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is useful as a runtime guard or when you want to load tool metadata for inspection without exposing the tools for invocation. + use_progressive_disclosure: When ``True``, expose discovery and loader tools plus + ``always_load`` tools initially, and let the model load and unload other allowed MCP tools on demand. + always_load: MCP tool names to keep visible from the start when progressive disclosure + is enabled. Names use the same safe matching rules as ``allowed_tools``; unmatched + entries are ignored. additional_properties: Additional properties. terminate_on_close: Close the transport when the MCP client is terminated. client: The chat client to use for sampling. @@ -2713,6 +3008,8 @@ def __init__( description=description, approval_mode=approval_mode, allowed_tools=allowed_tools, + use_progressive_disclosure=use_progressive_disclosure, + always_load=always_load, tool_name_prefix=tool_name_prefix, additional_properties=additional_properties, session=session, @@ -2850,6 +3147,8 @@ def __init__( description: str | None = None, approval_mode: (Literal["always_require", "never_require"] | MCPSpecificApproval | None) = None, allowed_tools: Collection[str] | None = None, + use_progressive_disclosure: bool = False, + always_load: Collection[str] | None = None, client: SupportsChatGetResponse | None = None, sampling_approval_callback: SamplingApprovalCallback | None = None, sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS, @@ -2903,6 +3202,11 @@ def __init__( disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is useful as a runtime guard or when you want to load tool metadata for inspection without exposing the tools for invocation. + use_progressive_disclosure: When ``True``, expose discovery and loader tools plus + ``always_load`` tools initially, and let the model load and unload other allowed MCP tools on demand. + always_load: MCP tool names to keep visible from the start when progressive disclosure + is enabled. Names use the same safe matching rules as ``allowed_tools``; unmatched + entries are ignored. additional_properties: Additional properties. client: The chat client to use for sampling. sampling_approval_callback: Optional gate run before each server-initiated @@ -2938,6 +3242,8 @@ def __init__( description=description, approval_mode=approval_mode, allowed_tools=allowed_tools, + use_progressive_disclosure=use_progressive_disclosure, + always_load=always_load, tool_name_prefix=tool_name_prefix, additional_properties=additional_properties, session=session, diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index c9322101f7..9b5e272edd 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -6,9 +6,10 @@ import logging import os import sys +import warnings from contextlib import _AsyncGeneratorContextManager # type: ignore from datetime import timedelta -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, Mock, patch import pytest @@ -21,11 +22,13 @@ Content, FunctionInvocationContext, FunctionMiddleware, + FunctionTool, MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool, Message, ) +from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning from agent_framework._mcp import ( MCPTool, _build_prefixed_mcp_name, @@ -56,6 +59,10 @@ def _mcp_result_to_text(result: str | list[Content]) -> str: _HELPER_MCP_TOOL = MCPTool(name="helper") # type: ignore[abstract] +def _reset_progressive_mcp_warning_state() -> None: + _WARNED_FEATURES.discard((ExperimentalWarning, ExperimentalFeature.PROGRESSIVE_TOOLS.value)) + + # Helper function tests def test_normalize_mcp_name(): """Test MCP name normalization.""" @@ -1707,6 +1714,643 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: assert sorted(actual_names) == sorted(expected_names) +def test_mcp_transport_subclasses_accept_progressive_disclosure_options() -> None: + _reset_progressive_mcp_warning_state() + with pytest.warns(ExperimentalWarning, match="PROGRESSIVE_TOOLS"): + stdio = MCPStdioTool( + name="stdio", + command="python", + use_progressive_disclosure=True, + always_load=["search"], + ) + http = MCPStreamableHTTPTool( + name="http", + url="https://example.com/mcp", + use_progressive_disclosure=True, + always_load=["search"], + ) + websocket = MCPWebsocketTool( + name="ws", + url="wss://example.com/mcp", + use_progressive_disclosure=True, + always_load=["search"], + ) + + assert stdio.use_progressive_disclosure is True + assert http.use_progressive_disclosure is True + assert websocket.use_progressive_disclosure is True + assert stdio.always_load == ["search"] + assert http.always_load == ["search"] + assert websocket.always_load == ["search"] + + +def test_mcp_progressive_disclosure_requires_loading_tools() -> None: + with pytest.raises(ValueError, match="requires load_tools=True"): + MCPTool( # type: ignore[abstract] + name="test_server", + load_tools=False, + use_progressive_disclosure=True, + ) + + +def test_mcp_progressive_disclosure_warns_on_construction() -> None: + _reset_progressive_mcp_warning_state() + with pytest.warns(ExperimentalWarning, match=f"{ExperimentalFeature.PROGRESSIVE_TOOLS.value}"): + MCPTool(name="test_server", use_progressive_disclosure=True) # type: ignore[abstract] + assert (ExperimentalWarning, ExperimentalFeature.PROGRESSIVE_TOOLS.value) in _WARNED_FEATURES + + +def test_mcp_tool_base_constructor_preserves_positional_tool_name_prefix() -> None: + tool = MCPTool("test_server", "description", None, None, "prefix") # type: ignore[abstract] + + assert tool.tool_name_prefix == "prefix" + assert tool.use_progressive_disclosure is False + + +def _progressive_tool_list_page(*, tools: list[types.Tool] | None = None) -> types.ListToolsResult: + if tools is not None: + return types.ListToolsResult(tools=tools) + return types.ListToolsResult( + tools=[ + types.Tool( + name="tool_one", + description="First tool", + inputSchema={ + "type": "object", + "properties": {"param": {"type": "string"}}, + "required": ["param"], + }, + ), + types.Tool( + name="tool_two", + description="Second tool", + inputSchema={ + "type": "object", + "properties": {"value": {"type": "integer"}}, + "required": ["value"], + }, + ), + types.Tool( + name="secret_tool", + description="Secret tool", + inputSchema={ + "type": "object", + "properties": {"secret": {"type": "string"}}, + }, + ), + ] + ) + + +async def _load_progressive_test_server( + *, + allowed_tools: list[str] | None = None, + always_load: list[str] | None = None, + tool_name_prefix: str | None = None, + approval_mode: Any = None, + tools: list[types.Tool] | None = None, +) -> MCPTool: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ExperimentalWarning) + server = MCPTool( # type: ignore[abstract] + name="test_server", + allowed_tools=allowed_tools, + always_load=always_load, + tool_name_prefix=tool_name_prefix, + approval_mode=approval_mode, + use_progressive_disclosure=True, + ) + server.session = AsyncMock() + server.session.list_tools = AsyncMock(return_value=_progressive_tool_list_page(tools=tools)) + server.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="ok")]) + ) + await server.load_tools() + return server + + +async def test_mcp_progressive_disclosure_exposes_loaders_and_always_loaded_tools() -> None: + server = await _load_progressive_test_server(always_load=["tool_one", "missing_tool"]) + + assert [func.name for func in server.functions] == ["list_mcp_tools", "load_tool", "unload_tool", "tool_one"] + + +async def test_mcp_progressive_disclosure_filters_always_loaded_loader_name_collisions() -> None: + server = await _load_progressive_test_server( + always_load=["list_mcp_tools", "tool_one"], + tools=[ + types.Tool( + name="list_mcp_tools", + description="Remote tool whose local name collides with the list loader.", + inputSchema={"type": "object", "properties": {}}, + ), + types.Tool( + name="tool_one", + description="First tool", + inputSchema={"type": "object", "properties": {}}, + ), + ], + ) + + assert [func.name for func in server.functions] == ["list_mcp_tools", "load_tool", "unload_tool", "tool_one"] + + +async def test_mcp_progressive_disclosure_loader_names_honor_tool_name_prefix() -> None: + server = await _load_progressive_test_server(tool_name_prefix="github", always_load=["tool_one"]) + + assert [func.name for func in server.functions] == [ + "github_list_mcp_tools", + "github_load_tool", + "github_unload_tool", + "github_tool_one", + ] + assert server.functions[1].parameters()["properties"]["tool"] == { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + "description": "The MCP tool name, or MCP tool names, to load.", + } + assert server.functions[2].parameters()["properties"]["tool"] == { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + "description": "The MCP tool name, or MCP tool names, to unload.", + } + + +async def test_mcp_progressive_list_mcp_tools_only_reports_allowed_tools() -> None: + server = await _load_progressive_test_server(allowed_tools=["tool_one"], always_load=[]) + list_tool = server.functions[0] + context = FunctionInvocationContext(function=list_tool, arguments={}, tools=list(server.functions)) + + result = await list_tool.invoke(arguments={}, context=context, skip_parsing=True) + + assert [item["name"] for item in result] == ["tool_one"] + assert result[0]["remote_name"] == "tool_one" + assert result[0]["description"] == "First tool" + assert result[0]["parameters"]["properties"] == {"param": {"type": "string"}} + assert result[0]["loaded"] is False + assert result[0]["always_loaded"] is False + + +async def test_mcp_progressive_list_mcp_tools_skips_loader_name_collisions() -> None: + server = await _load_progressive_test_server( + tools=[ + types.Tool( + name="list_mcp_tools", + description="Remote tool whose local name collides with the list loader.", + inputSchema={"type": "object", "properties": {}}, + ), + types.Tool( + name="tool_one", + description="First tool", + inputSchema={"type": "object", "properties": {}}, + ), + ], + ) + list_tool = server.functions[0] + context = FunctionInvocationContext(function=list_tool, arguments={}, tools=list(server.functions)) + + result = await list_tool.invoke(arguments={}, context=context, skip_parsing=True) + + assert [item["name"] for item in result] == ["tool_one"] + + +async def test_mcp_progressive_list_mcp_tools_treats_empty_allowed_tools_as_no_tools() -> None: + server = await _load_progressive_test_server(allowed_tools=[], always_load=["tool_one"]) + list_tool = server.functions[0] + context = FunctionInvocationContext(function=list_tool, arguments={}, tools=list(server.functions)) + + result = await list_tool.invoke(arguments={}, context=context, skip_parsing=True) + + assert result == [] + assert [func.name for func in server.functions] == ["list_mcp_tools", "load_tool", "unload_tool"] + + +async def test_mcp_progressive_load_tool_adds_hidden_tool_to_live_tools() -> None: + server = await _load_progressive_test_server() + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "tool_two"}, + tools=list(server.functions), + ) + + result = await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + + assert result[0].text == "Loaded MCP tool 'tool_two'. It is available on the next model iteration." + assert context.tools is not None + assert [tool.name for tool in context.tools] == [ # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_two", + ] + + list_tool = server.functions[0] + list_result = await list_tool.invoke(arguments={}, context=context, skip_parsing=True) + + assert next(item for item in list_result if item["name"] == "tool_two")["loaded"] is True + assert [func.name for func in server.functions] == [ + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_two", + ] + + result = await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + + assert result[0].text == "MCP tool 'tool_two' is already available." + assert [tool.name for tool in context.tools] == [ # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_two", + ] + + +async def test_mcp_progressive_load_tool_adds_multiple_hidden_tools_to_live_tools() -> None: + server = await _load_progressive_test_server() + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": ["tool_one", "tool_two"]}, + tools=list(server.functions), + ) + + result = await load_tool.invoke(arguments={"tool": ["tool_one", "tool_two"]}, context=context) + + assert result[0].text == ( + "Loaded MCP tool 'tool_one'. It is available on the next model iteration.\n" + "Loaded MCP tool 'tool_two'. It is available on the next model iteration." + ) + assert context.tools is not None + assert [tool.name for tool in context.tools] == [ # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_one", + "tool_two", + ] + assert [func.name for func in server.functions] == [ + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_one", + "tool_two", + ] + + +async def test_mcp_progressive_load_tool_reports_mixed_batch_results() -> None: + server = await _load_progressive_test_server(always_load=["tool_one"]) + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": ["tool_one", "tool_two", "missing_tool"]}, + tools=list(server.functions), + ) + + result = await load_tool.invoke(arguments={"tool": ["tool_one", "tool_two", "missing_tool"]}, context=context) + + assert result[0].text == ( + "MCP tool 'tool_one' is already available.\n" + "Loaded MCP tool 'tool_two'. It is available on the next model iteration.\n" + "MCP tool 'missing_tool' is not available. Available tools: tool_one, tool_two, secret_tool." + ) + assert context.tools is not None + assert [tool.name for tool in context.tools] == [ # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_one", + "tool_two", + ] + + +async def test_mcp_progressive_unload_tool_removes_dynamically_loaded_tool() -> None: + server = await _load_progressive_test_server() + load_tool = server.functions[1] + unload_tool = server.functions[2] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "tool_two"}, + tools=list(server.functions), + ) + await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + unload_context = FunctionInvocationContext( + function=unload_tool, + arguments={"tool": "tool_two"}, + tools=context.tools, + ) + + result = await unload_tool.invoke(arguments={"tool": "tool_two"}, context=unload_context) + + assert result[0].text == "Unloaded MCP tool 'tool_two'. It will be removed on the next model iteration." + assert unload_context.tools is not None + assert [tool.name for tool in unload_context.tools] == [ # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + "list_mcp_tools", + "load_tool", + "unload_tool", + ] + assert [func.name for func in server.functions] == ["list_mcp_tools", "load_tool", "unload_tool"] + + +async def test_mcp_progressive_unload_tool_removes_multiple_dynamically_loaded_tools() -> None: + server = await _load_progressive_test_server() + load_tool = server.functions[1] + unload_tool = server.functions[2] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": ["tool_one", "tool_two"]}, + tools=list(server.functions), + ) + await load_tool.invoke(arguments={"tool": ["tool_one", "tool_two"]}, context=context) + unload_context = FunctionInvocationContext( + function=unload_tool, + arguments={"tool": ["tool_one", "tool_two"]}, + tools=context.tools, + ) + + result = await unload_tool.invoke(arguments={"tool": ["tool_one", "tool_two"]}, context=unload_context) + + assert result[0].text == ( + "Unloaded MCP tool 'tool_one'. It will be removed on the next model iteration.\n" + "Unloaded MCP tool 'tool_two'. It will be removed on the next model iteration." + ) + assert unload_context.tools is not None + assert [tool.name for tool in unload_context.tools] == [ # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + "list_mcp_tools", + "load_tool", + "unload_tool", + ] + assert [func.name for func in server.functions] == ["list_mcp_tools", "load_tool", "unload_tool"] + + +async def test_mcp_progressive_unload_tool_rejects_always_loaded_tool() -> None: + server = await _load_progressive_test_server(always_load=["tool_one"]) + unload_tool = server.functions[2] + context = FunctionInvocationContext( + function=unload_tool, + arguments={"tool": "tool_one"}, + tools=list(server.functions), + ) + + result = await unload_tool.invoke(arguments={"tool": "tool_one"}, context=context) + + assert result[0].text == "MCP tool 'tool_one' is configured in always_load and cannot be unloaded." + assert context.tools is not None + visible_tools = cast(list[FunctionTool], context.tools) + assert [tool.name for tool in visible_tools] == [ + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_one", + ] + + +async def test_mcp_progressive_unload_tool_reports_tool_not_loaded() -> None: + server = await _load_progressive_test_server() + unload_tool = server.functions[2] + context = FunctionInvocationContext( + function=unload_tool, + arguments={"tool": "tool_two"}, + tools=list(server.functions), + ) + + result = await unload_tool.invoke(arguments={"tool": "tool_two"}, context=context) + + assert result[0].text == "MCP tool 'tool_two' is not currently loaded." + + +async def test_mcp_progressive_unload_tool_reports_mixed_batch_results() -> None: + server = await _load_progressive_test_server(always_load=["tool_one"]) + load_tool = server.functions[1] + unload_tool = server.functions[2] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "tool_two"}, + tools=list(server.functions), + ) + await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + unload_context = FunctionInvocationContext( + function=unload_tool, + arguments={"tool": ["tool_one", "tool_two", "secret_tool", "missing_tool"]}, + tools=context.tools, + ) + + result = await unload_tool.invoke( + arguments={"tool": ["tool_one", "tool_two", "secret_tool", "missing_tool"]}, + context=unload_context, + ) + + assert result[0].text == ( + "MCP tool 'tool_one' is configured in always_load and cannot be unloaded.\n" + "Unloaded MCP tool 'tool_two'. It will be removed on the next model iteration.\n" + "MCP tool 'secret_tool' is not currently loaded.\n" + "MCP tool 'missing_tool' is not available. Available tools: tool_one, tool_two, secret_tool." + ) + assert unload_context.tools is not None + visible_tools = cast(list[FunctionTool], unload_context.tools) + assert [tool.name for tool in visible_tools] == [ + "list_mcp_tools", + "load_tool", + "unload_tool", + "tool_one", + ] + assert [func.name for func in server.functions] == ["list_mcp_tools", "load_tool", "unload_tool", "tool_one"] + + +async def test_mcp_progressive_load_and_unload_suppress_experimental_context_warnings() -> None: + server = await _load_progressive_test_server() + load_tool = server.functions[1] + unload_tool = server.functions[2] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "tool_two"}, + tools=list(server.functions), + ) + + with warnings.catch_warnings(): + warnings.simplefilter("error", ExperimentalWarning) + await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + unload_context = FunctionInvocationContext( + function=unload_tool, + arguments={"tool": "tool_two"}, + tools=context.tools, + ) + await unload_tool.invoke(arguments={"tool": "tool_two"}, context=unload_context) + + +async def test_mcp_progressive_load_tool_rejects_filtered_tool() -> None: + server = await _load_progressive_test_server(allowed_tools=["tool_one"]) + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "tool_two"}, + tools=list(server.functions), + ) + + result = await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + + assert result[0].text == "MCP tool 'tool_two' is not available. Available tools: tool_one." + + +async def test_mcp_progressive_load_tool_rejects_missing_live_tool_list() -> None: + server = await _load_progressive_test_server() + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "tool_two"}, + ) + + with pytest.raises(ToolExecutionException, match="inside an agent function-calling run"): + await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + + +async def test_mcp_progressive_load_tool_rejects_loader_name_collision() -> None: + server = await _load_progressive_test_server( + tools=[ + types.Tool( + name="load_tool", + description="Remote tool whose local name collides with the load loader.", + inputSchema={"type": "object", "properties": {}}, + ), + ], + ) + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "load_tool"}, + tools=list(server.functions), + ) + + result = await load_tool.invoke(arguments={"tool": "load_tool"}, context=context) + + assert result[0].text is not None + assert "Set tool_name_prefix" in result[0].text + + +async def test_mcp_progressive_resolve_rejects_ambiguous_tool_name() -> None: + def _tool() -> None: + return None + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ExperimentalWarning) + server = MCPTool(name="test_server", use_progressive_disclosure=True) # type: ignore[abstract] + server._functions = [ + FunctionTool( + func=_tool, + name="remote_one", + additional_properties={"_mcp_remote_name": "shared", "_mcp_normalized_name": "remote_one"}, + ), + FunctionTool( + func=_tool, + name="remote_two", + additional_properties={"_mcp_remote_name": "shared", "_mcp_normalized_name": "remote_two"}, + ), + ] + + with pytest.raises(ToolExecutionException, match="ambiguous"): + server._resolve_progressive_function("shared") + + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "shared"}, + tools=list(server.functions), + ) + + result = await load_tool.invoke(arguments={"tool": "shared"}, context=context) + + assert result[0].text == "MCP tool name 'shared' is ambiguous." + + +async def test_mcp_progressive_loaded_tool_preserves_remote_approval_mode() -> None: + server = await _load_progressive_test_server(approval_mode={"always_require_approval": ["tool_two"]}) + load_tool = server.functions[1] + context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "tool_two"}, + tools=list(server.functions), + ) + + await load_tool.invoke(arguments={"tool": "tool_two"}, context=context) + + assert context.tools is not None + loaded_tool = cast( + FunctionTool, + next( + tool + for tool in context.tools + if tool.name == "tool_two" # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + ), + ) + assert loaded_tool.approval_mode == "always_require" + + +async def test_mcp_progressive_loaded_http_tool_preserves_runtime_kwargs_for_headers() -> None: + provider_received: list[dict[str, Any]] = [] + + def provider(kwargs: dict[str, Any]) -> dict[str, str]: + provider_received.append(dict(kwargs)) + return {"X-Some-Token": kwargs.get("some_token", "")} + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ExperimentalWarning) + server = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + header_provider=provider, + use_progressive_disclosure=True, + ) + server.session = AsyncMock() + server.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={ + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + }, + ) + ] + ) + ) + server.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + ) + await server.load_tools() + load_tool = server.functions[1] + load_context = FunctionInvocationContext( + function=load_tool, + arguments={"tool": "greet"}, + tools=list(server.functions), + ) + + await load_tool.invoke(arguments={"tool": "greet"}, context=load_context) + + assert load_context.tools is not None + loaded_tool = cast( + FunctionTool, + next( + tool + for tool in load_context.tools + if tool.name == "greet" # type: ignore[union-attr] # ty: ignore[unresolved-attribute] + ), + ) + invoke_context = FunctionInvocationContext( + function=loaded_tool, + arguments={"name": "Alice"}, + kwargs={"some_token": "my-secret"}, + ) + result = await loaded_tool.invoke(arguments={"name": "Alice"}, context=invoke_context) + + assert result[0].text == "Hello!" + assert provider_received[0]["some_token"] == "my-secret" + call_args = server.session.call_tool.call_args # type: ignore[union-attr] + assert call_args.kwargs.get("arguments", {}).get("name") == "Alice" + assert "some_token" not in call_args.kwargs.get("arguments", {}) + + # Server implementation tests def test_local_mcp_stdio_tool_init(): """Test MCPStdioTool initialization.""" @@ -2356,7 +3000,7 @@ async def test_mcp_tool_sampling_callback_forwards_system_prompt(): async def test_mcp_tool_sampling_callback_forwards_tools(): """Test sampling callback converts MCP tools to FunctionTools and passes them in options.""" - from agent_framework import FunctionTool, Message + from agent_framework import Message tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve) diff --git a/python/samples/02-agents/mcp/README.md b/python/samples/02-agents/mcp/README.md index 53af7d31a8..e3fce52794 100644 --- a/python/samples/02-agents/mcp/README.md +++ b/python/samples/02-agents/mcp/README.md @@ -14,6 +14,7 @@ The Model Context Protocol (MCP) is an open standard for connecting AI agents to | **API Key Authentication** | [`mcp_api_key_auth.py`](mcp_api_key_auth.py) | Demonstrates API key authentication with MCP servers using `header_provider`, runtime invocation kwargs, and a command-line API key argument | | **GitHub Integration with PAT** | [`mcp_github_pat.py`](mcp_github_pat.py) | Demonstrates connecting to GitHub's MCP server using Personal Access Token (PAT) authentication | | **Long-Running Task** | [`mcp_long_running_task.py`](mcp_long_running_task.py) | Demonstrates transparent SEP-2663 long-running task handling for MCP tools that advertise `taskSupport=required`. Self-spawns a stdio MCP child server | +| **Progressive Disclosure** | [`mcp_progressive_disclosure.py`](mcp_progressive_disclosure.py) | Demonstrates `use_progressive_disclosure`, `always_load`, `allowed_tools`, and prefixed `list_mcp_tools` / `load_tool` / `unload_tool` names. `load_tool` and `unload_tool` can accept one tool name or multiple names. Self-spawns a stdio MCP child server | | **Sampling Approval** | [`mcp_sampling_approval.py`](mcp_sampling_approval.py) | Demonstrates gating server-initiated `sampling/createMessage` requests with a `sampling_approval_callback`, plus the `sampling_max_tokens` and `sampling_max_requests` guardrails. MCP sampling is denied by default | ## Prerequisites @@ -25,6 +26,8 @@ Most samples in this folder use OpenAI: Run `mcp_api_key_auth.py` with the MCP API key as the first command-line argument. +`mcp_progressive_disclosure.py` self-spawns its demo MCP stdio server; no separate MCP server setup is required. + For `mcp_github_pat.py`: - `GITHUB_PAT` - Your GitHub Personal Access Token (create at https://github.com/settings/tokens) diff --git a/python/samples/02-agents/mcp/mcp_progressive_disclosure.py b/python/samples/02-agents/mcp/mcp_progressive_disclosure.py new file mode 100644 index 0000000000..4f48e12910 --- /dev/null +++ b/python/samples/02-agents/mcp/mcp_progressive_disclosure.py @@ -0,0 +1,154 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import sys +from typing import Any + +from agent_framework import Agent, MCPStdioTool +from agent_framework.openai import OpenAIChatClient +from dotenv import load_dotenv + +__doc__ = """ +MCP Progressive Disclosure Example + +This sample demonstrates how to connect an agent to a large MCP server without +frontloading every remote tool schema into the model prompt. + +The sample starts a tiny local MCP stdio server in a child process. The server +advertises three tools: + +1. ``get_server_status`` — always visible to the model. +2. ``search_docs`` — allowed, but hidden until the model calls ``docs_load_tool`` and removable with + ``docs_unload_tool`` when it is no longer useful. +3. ``internal_admin_report`` — not listed in ``allowed_tools``, so the model never + sees it in ``docs_list_mcp_tools`` and cannot load it. + +The ``MCPStdioTool`` is configured with: + +1. ``use_progressive_disclosure=True`` to enable loader tools. +2. ``always_load=["get_server_status"]`` to keep one cheap tool visible up front. +3. ``allowed_tools=[...]`` to define the only remote tools the model may discover + or load. +4. ``tool_name_prefix="docs"`` so multiple MCP servers can expose their own + ``docs_list_mcp_tools`` / ``docs_load_tool`` / ``docs_unload_tool`` names without collisions. + ``docs_load_tool`` and ``docs_unload_tool`` accept either one tool name or a list of tool names. + +Sample output: +User: Explain how progressive MCP tool disclosure works. First inspect the MCP +tools you can load, then load the docs search tool, use it, and unload it. +Agent: Progressive disclosure starts with a small set of tools. I listed the +available MCP tools, loaded docs_search_docs, and used it to find that hidden +MCP tools become available on the next function-calling iteration. +""" + + +load_dotenv() + + +async def _run_server() -> None: + """Run a minimal stdio MCP server with visible, loadable, and filtered tools.""" + import mcp.types as types + from mcp.server.lowlevel import Server + from mcp.server.stdio import stdio_server + + server: Server[Any, Any] = Server("mcp-progressive-disclosure-demo") + + @server.list_tools() + async def _list_tools() -> list[types.Tool]: # pyright: ignore[reportUnusedFunction] + return [ + types.Tool( + name="get_server_status", + description="Return the health of the demo MCP server.", + inputSchema={"type": "object", "properties": {}}, + ), + types.Tool( + name="search_docs", + description="Search short documentation snippets about MCP progressive disclosure.", + inputSchema={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The documentation search query.", + } + }, + "required": ["query"], + }, + ), + types.Tool( + name="internal_admin_report", + description="Internal server details that are intentionally filtered out by allowed_tools.", + inputSchema={"type": "object", "properties": {}}, + ), + ] + + @server.call_tool() + async def _call_tool(name: str, arguments: dict[str, Any]) -> types.CallToolResult: # pyright: ignore[reportUnusedFunction] + if name == "get_server_status": + text = "The demo MCP server is healthy. Use search_docs for progressive disclosure details." + elif name == "search_docs": + query = str(arguments.get("query", "")).strip() or "progressive disclosure" + text = ( + f"Search results for '{query}': In progressive MCP disclosure, the agent starts with " + "list/load/unload tools and selected always-loaded tools. Calling load_tool adds an allowed " + "remote MCP tool to the live tool list for the next model iteration, and unload_tool removes it." + ) + elif name == "internal_admin_report": + text = "This tool should not be discoverable because it is excluded by allowed_tools." + else: + text = f"Unknown tool: {name}" + return types.CallToolResult(content=[types.TextContent(type="text", text=text)]) + + async with stdio_server() as (read_stream, write_stream): + await server.run(read_stream, write_stream, server.create_initialization_options()) + + +async def _run_client() -> None: + """Run an agent that progressively discovers and loads MCP tools.""" + # 1. Create the MCP tool. Only get_server_status is visible at first; search_docs + # is discoverable through docs_list_mcp_tools, loadable through docs_load_tool, + # and unloadable through docs_unload_tool. + mcp_tool = MCPStdioTool( + name="DocsMCP", + description="Demo MCP server with progressively loaded documentation tools.", + command=sys.executable, + args=[__file__, "--server"], + allowed_tools=["get_server_status", "search_docs"], + use_progressive_disclosure=True, + always_load=["get_server_status"], + tool_name_prefix="docs", + ) + + # 2. Create an agent with the progressive MCP tool. + async with Agent( + client=OpenAIChatClient(), + name="ProgressiveMCPAgent", + instructions=( + "You are a helpful assistant. To answer documentation questions, first call " + "docs_list_mcp_tools to see which MCP tools are available. If you need a " + "hidden tool, call docs_load_tool with that tool's remote name, then call " + "the newly available prefixed tool on the next iteration. When the hidden " + "tool is no longer needed, call docs_unload_tool. Do not invent tools that are not listed." + ), + tools=mcp_tool, + ) as agent: + # 3. Ask a question that requires loading a hidden MCP tool. + prompt = ( + "Explain how progressive MCP tool disclosure works. First inspect the MCP " + "tools you can load, then load the docs search tool, use it, and unload it." + ) + print(f"User: {prompt}") + response = await agent.run(prompt) + print(f"Agent: {response.text}") + + +async def main() -> None: + """Run either the MCP server branch or the agent client branch.""" + if "--server" in sys.argv: + await _run_server() + else: + await _run_client() + + +if __name__ == "__main__": + asyncio.run(main())