diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index f9438a887a0..8c72566c7a3 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -161,8 +161,9 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **`MCPTool`** - Base wrapper that owns the MCP `ClientSession` and exposes the remote server's tools as `FunctionTool`s. - **`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. **The declared half comes from the server's advertised schema**, and runtime kwargs (`FunctionInvocationContext.kwargs`, seeded from `function_invocation_kwargs`) are merged with the model-supplied arguments upstream in `_call_tool_with_runtime_kwargs`, so provenance is gone by the time the filter runs. A runtime kwarg is therefore forwarded whenever the server declares a property of that name, without the model mentioning it — the server, not the caller, decides which runtime kwarg names it receives. A tool that declares no usable `properties` (including schemas with `additionalProperties: true`) forwards only the configured extras. `_MCP_FRAMEWORK_DENYLIST` is a narrow safety net covering only non-serializable framework objects a server *declares* in its schema (those are dropped); it does not generalize to arbitrary caller-chosen names, and explicit extras 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. +- **`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. Configured allow/approval names must identify at most one raw remote name across loaded tools and prompts: a raw name that overlaps another tool's prefixed alias raises `ToolExecutionException`. Discovery validates all pages before publishing new functions; a failed reload retains the previous functions and metadata. The allowlist is also revalidated when exposing functions so runtime changes cannot select an ambiguous 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. +- **Tool discovery refresh** - Validate the complete current tool snapshot alongside retained non-tool functions, not stale tool wrappers. Successful refreshes atomically replace the previous tool set and metadata, including empty snapshots, while preserving prompts, custom functions, and unchanged tool wrappers (including caller customizations). Removed or replaced tool identities lose their persisted progressive-loaded state; failed refreshes preserve the previous state. - **`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 — but note this constrains the *model*, not the *server*, which still widens the effective allowlist through its schema. 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. - **MCP HTTP header request scoping** - `static_headers` supplies fixed, origin-scoped headers without serializing concurrent calls; `header_provider` resolves dynamic per-call headers under the instance call lock. Fixed and dynamic headers form the session's complete effective identity (case-insensitive names, case-sensitive values), with dynamic values overriding fixed values of the same name. Agent runs reconcile a connected tool before copying its functions; generated tool and prompt calls perform the same check at invocation. Connected run preparation defers when a provider needs model-supplied values that are not available until invocation, where resolution remains strict. A different identity waits for in-flight discovery, reconnects framework-created sessions without loading, then refreshes session-derived discovery under the same lock. Because a caller-supplied session's established identity is unknown and the wrapper cannot reconnect it, dynamic header resolution on that wrapper is rejected even while disconnected. Ambient requests retain the bound set. With a shared `http_client`, processing stays scoped to the originating `MCPStreamableHTTPTool`, and every injected header is stripped from cross-origin redirects. The session exit stack removes its request hook after transport shutdown, including failed initialization or discovery, and closes framework-created HTTP clients. Failed discovery resets the connection and discovery flags and rolls back partial function/metadata additions. Framework-created sessions are discarded; constructor-supplied sessions remain caller-owned and reusable across cleanup, close, and reset. - **MCP lifecycle caller cancellation** - The lifecycle owner skips cancelled queued connect requests and waits for the caller to acknowledge successful setup. If the caller cancels before accepting a newly established connection, the owner tears it down before processing the next request. Cancelling a close waiter does not interrupt teardown, and cancelling a redundant connect does not discard a previously established session. diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 0dd82180792..2073141ad94 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -87,6 +87,9 @@ class MCPSpecificApproval(TypedDict, total=False): """Represents the specific approval mode for an MCP tool. When using this mode, the user must specify which tools always or never require approval. + Names match raw remote names or unambiguous prefixed names when the raw name already + matches its normalized form. Normalized-only aliases do not match. Discovery raises + ``ToolExecutionException`` if a configured name matches multiple raw remote names. Attributes: always_require_approval: A sequence of tool names that always require approval. @@ -97,7 +100,6 @@ class MCPSpecificApproval(TypedDict, total=False): never_require_approval: Collection[str] | None - MCPToolResultContentMode = Literal[ "structured_first", "content_first", @@ -121,6 +123,7 @@ class MCPSpecificApproval(TypedDict, total=False): """ _MCP_REMOTE_NAME_KEY = "_mcp_remote_name" +_MCP_IS_TOOL_KEY = "_mcp_is_tool" _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name" _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY = "_mcp_tool_result_host_payload" _MCP_PROGRESSIVE_LIST_TOOL_NAME = "list_mcp_tools" @@ -909,6 +912,7 @@ def __init__( A non-empty collection exposes only the raw remote tools whose names appear in it. For compatibility, the prefixed local function name is also accepted when the raw remote name already matches its normalized form; normalized aliases do not authorize a different raw remote tool. + A configured name matching multiple raw remote names raises ``ToolExecutionException``. An empty collection (``[]``) exposes no tools — if you simply want to disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is useful as a runtime guard or when you want to load tool metadata for @@ -986,8 +990,7 @@ def __init__( } if tool_result_content not in allowed_modes: raise ValueError( - f"tool_result_content must be one of {sorted(allowed_modes)}, " - f"got {tool_result_content!r}." + f"tool_result_content must be one of {sorted(allowed_modes)}, got {tool_result_content!r}." ) self.name = name self.description = description or "" @@ -1373,6 +1376,7 @@ def functions(self) -> list[FunctionTool]: def _filtered_functions(self) -> list[FunctionTool]: """Return loaded MCP functions after applying ``allowed_tools``.""" + self._validate_config_names(self._functions) if self.allowed_tools is None: return self._functions allowed_names = set(self.allowed_tools) @@ -1382,6 +1386,38 @@ def _filtered_functions(self) -> list[FunctionTool]: filtered_functions.append(func) return filtered_functions + def _validate_config_names(self, functions: Sequence[FunctionTool]) -> None: + """Reject configured names that identify more than one raw remote name.""" + configured_names = set(self.allowed_tools or ()) + if isinstance(self.approval_mode, dict): + configured_names.update(self.approval_mode.get("always_require_approval") or ()) + configured_names.update(self.approval_mode.get("never_require_approval") or ()) + if not configured_names: + return + + remote_by_config_name: dict[str, str] = {} + for func in functions: + 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): + continue + for name in _mcp_config_candidate_names( + local_name=func.name, + normalized_name=normalized_name, + remote_name=remote_name, + ): + if name not in configured_names: + continue + existing_remote = remote_by_config_name.get(name) + if existing_remote is not None and existing_remote != remote_name: + raise ToolExecutionException( + f"MCP configuration name {name!r} is ambiguous: it matches raw remote names " + f"{existing_remote!r} and {remote_name!r}. " + "Use an unambiguous name or a different tool_name_prefix." + ) + remote_by_config_name[name] = remote_name + 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: @@ -2341,7 +2377,8 @@ async def load_prompts(self) -> None: them into FunctionTool instances. Handles pagination automatically. Raises: - ToolExecutionException: If the MCP server is not connected. + ToolExecutionException: If the MCP server is not connected or a configured + allow/approval name matches multiple raw remote names. """ async with self._function_load_lock: await self._load_prompts_locked() @@ -2356,6 +2393,7 @@ async def _load_prompts_locked(self) -> None: # Track existing function names to prevent duplicates existing_names = {func.name for func in self._functions} + new_functions: list[FunctionTool] = [] params: types.PaginatedRequestParams | None = None while True: @@ -2419,7 +2457,7 @@ async def _load_prompts_locked(self) -> None: _MCP_NORMALIZED_NAME_KEY: normalized_name, }, ) - self._functions.append(func) + new_functions.append(func) existing_names.add(local_name) # Check if there are more pages @@ -2427,6 +2465,9 @@ async def _load_prompts_locked(self) -> None: break params = types.PaginatedRequestParams(cursor=prompt_list.nextCursor) + self._validate_config_names([*self._functions, *new_functions]) + self._functions.extend(new_functions) + async def load_tools(self) -> None: """Load tools from the MCP server. @@ -2434,7 +2475,9 @@ async def load_tools(self) -> None: them into FunctionTool instances. Handles pagination automatically. Raises: - ToolExecutionException: If the MCP server is not connected. + ToolExecutionException: If the MCP server is not connected, multiple tools + map to the same local name, or a configured allow/approval name matches + multiple raw remote names. """ async with self._function_load_lock: await self._load_tools_locked() @@ -2447,9 +2490,15 @@ async def _load_tools_locked(self) -> None: logger.debug("Skipping MCP tool loading because the server did not advertise tools support.") return - # Track existing function names to prevent duplicates + # Previous tools are reusable; seed duplicate detection only with retained non-tool functions. + existing_tools: dict[str, FunctionTool] = {} existing_remote_by_local: dict[str, str] = {} + new_functions: list[FunctionTool] = [] + reused_tool_names: set[str] = set() for func in self._functions: + if (func.additional_properties or {}).get(_MCP_IS_TOOL_KEY): + existing_tools[func.name] = func + continue remote_name = (func.additional_properties or {}).get(_MCP_REMOTE_NAME_KEY) if isinstance(remote_name, str): existing_remote_by_local[func.name] = remote_name @@ -2531,6 +2580,14 @@ async def _load_tools_locked(self) -> None: existing_remote_by_local[local_name] = tool.name + existing_tool = existing_tools.get(local_name) + if ( + existing_tool is not None + and (existing_tool.additional_properties or {}).get(_MCP_REMOTE_NAME_KEY) == tool.name + ): + reused_tool_names.add(local_name) + continue + approval_mode = self._determine_approval_mode( *_mcp_config_candidate_names( local_name=local_name, @@ -2549,18 +2606,28 @@ async def _load_tools_locked(self) -> None: additional_properties={ _MCP_REMOTE_NAME_KEY: tool.name, _MCP_NORMALIZED_NAME_KEY: normalized_name, + _MCP_IS_TOOL_KEY: True, }, ) - self._functions.append(func) + new_functions.append(func) # Check if there are more pages if not tool_list.nextCursor: break params = types.PaginatedRequestParams(cursor=tool_list.nextCursor) + current_functions = [ + func + for func in self._functions + if not (func.additional_properties or {}).get(_MCP_IS_TOOL_KEY) or func.name in reused_tool_names + ] + current_functions.extend(new_functions) + self._validate_config_names(current_functions) + self._functions[:] = current_functions self._tool_call_meta_by_name = tool_call_meta_by_name self._tool_task_support_by_name = tool_task_support_by_name self._tool_param_names_by_name = tool_param_names_by_name + self._progressive_loaded_tool_names.difference_update(existing_tools.keys() - reused_tool_names) async def _cancel_pending_reload_tasks(self) -> None: """Cancel session-bound discovery reloads and wait for them to finish.""" @@ -3477,6 +3544,8 @@ def __init__( allowed_tools: Optional allow-list of MCP tool names to expose as functions. ``None`` (the default) exposes every tool advertised by the MCP server. A non-empty collection exposes only the tools whose names appear in it. + Names match raw remote names or unambiguous prefixed names when normalization + leaves the raw name unchanged. Ambiguous allow/approval names raise ``ToolExecutionException``. An empty collection (``[]``) exposes no tools — if you simply want to disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is useful as a runtime guard or when you want to load tool metadata for @@ -3685,6 +3754,8 @@ def __init__( allowed_tools: Optional allow-list of MCP tool names to expose as functions. ``None`` (the default) exposes every tool advertised by the MCP server. A non-empty collection exposes only the tools whose names appear in it. + Names match raw remote names or unambiguous prefixed names when normalization + leaves the raw name unchanged. Ambiguous allow/approval names raise ``ToolExecutionException``. An empty collection (``[]``) exposes no tools — if you simply want to disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is useful as a runtime guard or when you want to load tool metadata for @@ -4279,6 +4350,8 @@ def __init__( allowed_tools: Optional allow-list of MCP tool names to expose as functions. ``None`` (the default) exposes every tool advertised by the MCP server. A non-empty collection exposes only the tools whose names appear in it. + Names match raw remote names or unambiguous prefixed names when normalization + leaves the raw name unchanged. Ambiguous allow/approval names raise ``ToolExecutionException``. An empty collection (``[]``) exposes no tools — if you simply want to disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is useful as a runtime guard or when you want to load tool metadata for diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 8a5ccf18a30..c8f547086f4 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -38,6 +38,7 @@ from agent_framework._mcp import ( _MCP_HEADER_OWNER_EXTENSION, _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, + MCPSpecificApproval, MCPTool, _build_prefixed_mcp_name, _describe_error, @@ -171,13 +172,14 @@ def test_mcp_transport_subclasses_accept_tool_name_prefix() -> None: ) -async def test_load_tools_with_tool_name_prefix_preserves_matching_configuration(): +@pytest.mark.parametrize("configured_name", ["search_docs", "docs_search_docs"]) +async def test_load_tools_with_tool_name_prefix_preserves_matching_configuration(configured_name: str): """Prefixed MCP tool names should still honor unprefixed allow/approval configuration.""" tool = MCPTool( # type: ignore[abstract] name="docs", tool_name_prefix="docs", - allowed_tools=["search_docs"], - approval_mode={"always_require_approval": ["search_docs"]}, + allowed_tools=[configured_name], + approval_mode={"always_require_approval": [configured_name]}, ) mock_session = AsyncMock() @@ -202,6 +204,343 @@ async def test_load_tools_with_tool_name_prefix_preserves_matching_configuration assert tool.functions[0].approval_mode == "always_require" +@pytest.mark.parametrize( + "allowed_tools,approval_mode", + [ + (["docs_search"], None), + (None, {"always_require_approval": ["docs_search"]}), + (None, {"never_require_approval": ["docs_search"]}), + ], + ids=["allowed", "always_require", "never_require"], +) +@pytest.mark.parametrize("remote_names", [("search", "docs_search"), ("docs_search", "search")]) +@pytest.mark.parametrize("paginated", [False, True]) +@pytest.mark.parametrize("progressive", [False, True]) +async def test_load_tools_rejects_ambiguous_policy_names( + allowed_tools: list[str] | None, + approval_mode: MCPSpecificApproval | None, + remote_names: tuple[str, str], + paginated: bool, + progressive: bool, +) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ExperimentalWarning) + tool = MCPTool( # type: ignore[abstract] + name="docs", + tool_name_prefix="docs", + allowed_tools=allowed_tools, + approval_mode=approval_mode, + use_progressive_disclosure=progressive, + ) + advertised_tools = [ + types.Tool(name=name, inputSchema={"type": "object", "properties": {}}) for name in remote_names + ] + + async def list_tools(params: types.PaginatedRequestParams | None = None) -> types.ListToolsResult: + assert tool._functions == [] + if paginated: + if params is None: + return types.ListToolsResult(tools=advertised_tools[:1], nextCursor="second") + assert params.cursor == "second" + return types.ListToolsResult(tools=advertised_tools[1:]) + return types.ListToolsResult(tools=advertised_tools) + + tool.session = AsyncMock() + tool.session.list_tools = AsyncMock(side_effect=list_tools) + + with pytest.raises(ToolExecutionException, match="configuration name 'docs_search' is ambiguous"): + await tool.load_tools() + + assert tool._functions == [] + assert tool._tool_param_names_by_name == {} + assert [function.name for function in tool.functions] == ( + ["docs_list_mcp_tools", "docs_load_tool", "docs_unload_tool"] if progressive else [] + ) + + +@pytest.mark.parametrize( + "allowed_tools,approval_mode", + [ + (["docs_search"], None), + (None, {"always_require_approval": ["docs_search"]}), + (None, {"never_require_approval": ["docs_search"]}), + ], + ids=["allowed", "always_require", "never_require"], +) +@pytest.mark.parametrize("remote_names", [("search", "docs_search"), ("docs_search", "search")]) +async def test_ambiguous_policy_reload_preserves_previous_discovery( + allowed_tools: list[str] | None, + approval_mode: MCPSpecificApproval | None, + remote_names: tuple[str, str], + caplog: pytest.LogCaptureFixture, +) -> None: + tool = MCPTool( # type: ignore[abstract] + name="docs", tool_name_prefix="docs", allowed_tools=allowed_tools, approval_mode=approval_mode + ) + tool.session = AsyncMock() + original = types.Tool( + name=remote_names[0], + inputSchema={"type": "object", "properties": {"query": {"type": "string"}}}, + _meta={"original": True}, + ) + tool.session.list_tools = AsyncMock(return_value=types.ListToolsResult(tools=[original])) + await tool.load_tools() + original_functions = list(tool._functions) + original_meta = dict(tool._tool_call_meta_by_name) + original_params = dict(tool._tool_param_names_by_name) + original_tasks = dict(tool._tool_task_support_by_name) + + async def list_tools(params: types.PaginatedRequestParams | None = None) -> types.ListToolsResult: + assert tool._functions == original_functions + if params is None: + return types.ListToolsResult( + tools=[original, types.Tool(name="extra", inputSchema={"type": "object"})], + nextCursor="second", + ) + return types.ListToolsResult(tools=[types.Tool(name=remote_names[1], inputSchema={"type": "object"})]) + + tool.session.list_tools = AsyncMock(side_effect=list_tools) + with caplog.at_level(logging.WARNING, logger=logger.name): + await tool.message_handler(types.ServerNotification(types.ToolListChangedNotification())) + await asyncio.gather(*tool._pending_reload_tasks) + + assert "configuration name 'docs_search' is ambiguous" in caplog.text + assert tool._functions == original_functions + assert tool.functions == original_functions + assert tool._tool_call_meta_by_name == original_meta + assert tool._tool_param_names_by_name == original_params + assert tool._tool_task_support_by_name == original_tasks + + +@pytest.mark.parametrize( + "allowed_tools,approval_mode,expected_approval", + [ + (["docs_search"], None, "never_require"), + (None, {"always_require_approval": ["docs_search"]}, "always_require"), + (None, {"never_require_approval": ["docs_search"]}, "never_require"), + ], + ids=["allowed", "always_require", "never_require"], +) +@pytest.mark.parametrize("remote_names", [("search", "docs_search"), ("docs_search", "search")]) +async def test_tool_refresh_accepts_unambiguous_rename( + allowed_tools: list[str] | None, + approval_mode: MCPSpecificApproval | None, + expected_approval: str, + remote_names: tuple[str, str], +) -> None: + tool = MCPTool( # type: ignore[abstract] + name="docs", tool_name_prefix="docs", allowed_tools=allowed_tools, approval_mode=approval_mode + ) + tool.session = AsyncMock() + tool.session.list_tools = AsyncMock( + side_effect=[ + types.ListToolsResult(tools=[types.Tool(name=name, inputSchema={"type": "object"})]) + for name in remote_names + ] + ) + await tool.load_tools() + await tool.message_handler(types.ServerNotification(types.ToolListChangedNotification())) + await asyncio.gather(*tool._pending_reload_tasks) + + assert [function.name for function in tool.functions] == [f"docs_{remote_names[1]}"] + assert tool.functions[0].approval_mode == expected_approval + tool.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="ok")]) + ) + await tool.functions[0].invoke(arguments={}) + assert tool.session.call_tool.call_args.args[0] == remote_names[1] + + +@pytest.mark.parametrize("empty_snapshot", [False, True]) +async def test_tool_refresh_replaces_snapshot_and_preserves_other_functions(empty_snapshot: bool) -> None: + tool = MCPTool(name="docs", tool_name_prefix="docs") # type: ignore[abstract] + tool.session = AsyncMock() + keep = types.Tool( + name="keep", + inputSchema={"type": "object", "properties": {"query": {"type": "string"}}}, + _meta={"version": 1}, + ) + removed = types.Tool( + name="removed", + inputSchema={"type": "object", "properties": {}}, + _meta={"version": 1}, + execution=types.ToolExecution(taskSupport="required"), + ) + tool.session.list_tools = AsyncMock(return_value=types.ListToolsResult(tools=[keep, removed])) + await tool.load_tools() + kept_function = tool._functions[0] + parser = Mock(return_value="custom result") + kept_function.result_parser = parser + tool.session.list_prompts = AsyncMock(return_value=types.ListPromptsResult(prompts=[types.Prompt(name="summary")])) + await tool.load_prompts() + prompt_function = tool._functions[-1] + custom_function = FunctionTool(name="custom", func=lambda: "custom") + tool._functions.append(custom_function) + original_list = tool._functions + original_functions = list(original_list) + original_meta = dict(tool._tool_call_meta_by_name) + keep.meta = {"version": 2} + + async def list_tools(params: types.PaginatedRequestParams | None = None) -> types.ListToolsResult: + assert tool._functions == original_functions + assert tool._tool_call_meta_by_name == original_meta + if params is None: + return types.ListToolsResult(tools=[], nextCursor="second") + assert params.cursor == "second" + return types.ListToolsResult(tools=[] if empty_snapshot else [keep]) + + tool.session.list_tools = AsyncMock(side_effect=list_tools) + await tool.load_tools() + + assert tool._functions is original_list + assert {function.name for function in tool._functions} == ( + {"docs_summary", "custom"} if empty_snapshot else {"docs_keep", "docs_summary", "custom"} + ) + assert next(function for function in tool._functions if function.name == "docs_summary") is prompt_function + assert next(function for function in tool._functions if function.name == "custom") is custom_function + if not empty_snapshot: + assert next(function for function in tool._functions if function.name == "docs_keep") is kept_function + assert kept_function.result_parser is parser + assert tool._tool_call_meta_by_name == ({} if empty_snapshot else {"keep": {"version": 2}}) + assert tool._tool_param_names_by_name == ({} if empty_snapshot else {"keep": {"query"}}) + assert tool._tool_task_support_by_name == {} + + +async def test_tool_refresh_preserves_prompt_when_server_advertises_same_raw_name() -> None: + tool = MCPTool(name="docs") # type: ignore[abstract] + tool.session = AsyncMock() + tool.session.list_prompts = AsyncMock(return_value=types.ListPromptsResult(prompts=[types.Prompt(name="summary")])) + await tool.load_prompts() + prompt_function = tool._functions[0] + tool.session.list_tools = AsyncMock( + side_effect=[ + types.ListToolsResult(tools=[types.Tool(name="summary", inputSchema={"type": "object"})]), + types.ListToolsResult(tools=[]), + ] + ) + await tool.load_tools() + await tool.load_tools() + + assert tool._functions == [prompt_function] + + +@pytest.mark.parametrize("replacement_name", [None, "search-docs"]) +async def test_tool_refresh_forgets_removed_progressive_tools(replacement_name: str | None) -> None: + tool = await _load_progressive_test_server( + tool_name_prefix="docs", + tools=[types.Tool(name=name, inputSchema={"type": "object"}) for name in ("search/docs", "keep")], + ) + loader = tool.functions[1] + context = FunctionInvocationContext(function=loader, arguments={}, tools=list(tool.functions)) + await loader.invoke(arguments={"tool": ["docs_search-docs", "keep"]}, context=context) + assert tool._progressive_loaded_tool_names == {"docs_search-docs", "docs_keep"} + kept_function = next(function for function in tool._functions if function.name == "docs_keep") + assert tool.session is not None + with patch.object( + tool.session, + "list_tools", + return_value=types.ListToolsResult( + tools=[ + types.Tool(name=name, inputSchema={"type": "object"}) + for name in (["keep", replacement_name] if replacement_name else ["keep"]) + ] + ), + ): + await tool.load_tools() + + assert tool._progressive_loaded_tool_names == {"docs_keep"} + assert tool.functions[-1] is kept_function + assert [function.name for function in tool.functions] == [ + "docs_list_mcp_tools", + "docs_load_tool", + "docs_unload_tool", + "docs_keep", + ] + assert [function.name for function in tool._functions] == ( + ["docs_keep", "docs_search-docs"] if replacement_name else ["docs_keep"] + ) + + +@pytest.mark.parametrize( + "allowed_tools,expected_names", + [ + (None, ["docs_search", "docs_docs_search"]), + ([], []), + (["search"], ["docs_search"]), + (["docs_docs_search"], ["docs_docs_search"]), + ], +) +async def test_overlapping_names_accept_unambiguous_policy( + allowed_tools: list[str] | None, expected_names: list[str] +) -> None: + tool = MCPTool( # type: ignore[abstract] + name="docs", + tool_name_prefix="docs", + allowed_tools=allowed_tools, + approval_mode={ + "always_require_approval": ["search"], + "never_require_approval": ["search", "docs_docs_search"], + }, + ) + tool.session = AsyncMock() + tool.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool(name=name, inputSchema={"type": "object", "properties": {}}) + for name in ("search", "docs_search") + ] + ) + ) + await tool.load_tools() + + assert [function.name for function in tool.functions] == expected_names + assert [function.approval_mode for function in tool._functions] == ["always_require", "never_require"] + + tool.allowed_tools = ["docs_search"] + with pytest.raises(ToolExecutionException, match="configuration name 'docs_search' is ambiguous"): + _ = tool.functions + + +@pytest.mark.parametrize("load_order", ["prompts", "tools_then_prompts", "prompts_then_tools"]) +async def test_prompts_reject_ambiguous_policy_names(load_order: str) -> None: + tool = MCPTool( # type: ignore[abstract] + name="docs", + tool_name_prefix="docs", + allowed_tools=["docs_search"], + ) + tool.session = AsyncMock() + tool.session.list_prompts = AsyncMock( + side_effect=[ + types.ListPromptsResult(prompts=[types.Prompt(name="search")], nextCursor="second"), + types.ListPromptsResult(prompts=[types.Prompt(name="docs_search")]), + ] + ) + if load_order == "prompts": + with pytest.raises(ToolExecutionException, match="configuration name 'docs_search' is ambiguous"): + await tool.load_prompts() + assert tool._functions == [] + return + + tool.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[types.Tool(name="search", inputSchema={"type": "object", "properties": {}})] + ) + ) + tool.session.list_prompts = AsyncMock( + return_value=types.ListPromptsResult(prompts=[types.Prompt(name="docs_search")]) + ) + first_load, second_load = ( + (tool.load_tools, tool.load_prompts) + if load_order == "tools_then_prompts" + else (tool.load_prompts, tool.load_tools) + ) + await first_load() + original_functions = list(tool._functions) + with pytest.raises(ToolExecutionException, match="configuration name 'docs_search' is ambiguous"): + await second_load() + assert tool._functions == original_functions + + async def test_allowed_tools_does_not_authorize_normalized_remote_name_collision() -> None: """A normalized/local allowlist match must not authorize a different raw remote tool.""" tool = MCPTool(name="test_server", allowed_tools=["delete-file"]) # type: ignore[abstract]