diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 56f68cf3c3b..89d32f9000d 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -515,6 +515,12 @@ async def delete(self, *args: Any, **kwargs: Any) -> Any: } +def _is_mcp_prompt_function(func: FunctionTool) -> bool: + """Return whether a function was created from an MCP prompt.""" + properties = func.additional_properties or {} + return isinstance(properties.get(_MCP_REMOTE_NAME_KEY), str) and not properties.get(_MCP_IS_TOOL_KEY) + + def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> dict[str, Any]: """Get the input model from an MCP prompt. @@ -2441,8 +2447,17 @@ async def _load_prompts_locked(self) -> None: logger.debug("Skipping MCP prompt loading because the server did not advertise prompts support.") return - # Track existing function names to prevent duplicates - existing_names = {func.name for func in self._functions} + # A complete prompts/list snapshot replaces previously loaded prompts; seed duplicate + # detection only with retained non-prompt functions (MCP tools and user-added functions). + existing_prompts: dict[str, FunctionTool] = {} + existing_names: set[str] = set() + for func in self._functions: + if _is_mcp_prompt_function(func): + existing_prompts[func.name] = func + else: + existing_names.add(func.name) + prompt_functions: list[FunctionTool] = [] + retained_prompt_names: set[str] = set() new_functions: list[FunctionTool] = [] params: types.PaginatedRequestParams | None = None @@ -2507,19 +2522,39 @@ async def _load_prompts_locked(self) -> None: _MCP_NORMALIZED_NAME_KEY: normalized_name, }, ) - new_functions.append(func) existing_names.add(local_name) + existing_prompt = existing_prompts.get(local_name) + if ( + existing_prompt is not None + and (existing_prompt.additional_properties or {}).get(_MCP_REMOTE_NAME_KEY) == prompt.name + ): + # Same remote prompt identity: keep its progressive-load state even if its metadata changed. + retained_prompt_names.add(local_name) + if ( + existing_prompt.description == func.description + and existing_prompt.parameters() == func.parameters() + ): + # Keep the unchanged prompt function so local customizations survive a refresh. + prompt_functions.append(existing_prompt) + continue + prompt_functions.append(func) + new_functions.append(func) # Check if there are more pages if not prompt_list.nextCursor: break params = types.PaginatedRequestParams(cursor=prompt_list.nextCursor) - self._validate_config_names([*self._functions, *new_functions]) + current_functions = [func for func in self._functions if not _is_mcp_prompt_function(func)] + current_functions.extend(prompt_functions) + self._validate_config_names(current_functions) if self._function_load_callback is not None: for function in new_functions: self._function_load_callback(function, None) - self._functions.extend(new_functions) + self._functions[:] = current_functions + # Only prompts that keep the same remote identity retain their progressive-load state; a + # different remote prompt that normalizes to the same local name must start unloaded. + self._progressive_loaded_tool_names.difference_update(existing_prompts.keys() - retained_prompt_names) async def load_tools(self) -> None: """Load tools from the MCP server. diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 367bdcdc7f5..47ca0cde917 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -672,6 +672,150 @@ async def test_load_prompts_with_tool_name_prefix() -> None: assert [function.name for function in tool._functions] == ["docs_summarize-docs"] +async def test_load_prompts_refresh_replaces_prompt_catalog() -> None: + """A complete prompts/list refresh removes, updates, and keeps prompts without touching other functions.""" + tool = MCPTool(name="docs") # type: ignore[abstract] # ty: ignore[call-non-callable] + tool.session = AsyncMock() + tool.session.list_tools = AsyncMock( + return_value=types.ListToolsResult(tools=[types.Tool(name="search", inputSchema={"type": "object"})]) + ) + await tool.load_tools() + tool_function = tool._functions[0] + tool.session.list_prompts = AsyncMock( + return_value=types.ListPromptsResult( + prompts=[ + types.Prompt(name="keep", description="Keep"), + types.Prompt(name="removed", description="Removed"), + types.Prompt(name="changed", description="Old"), + ] + ) + ) + await tool.load_prompts() + kept_prompt = next(function for function in tool._functions if function.name == "keep") + custom_function = FunctionTool(name="custom", func=lambda: "custom") + tool._functions.append(custom_function) + tool._progressive_loaded_tool_names.update({"keep", "removed"}) + + tool.session.list_prompts = AsyncMock( + side_effect=[ + types.ListPromptsResult(prompts=[types.Prompt(name="keep", description="Keep")], nextCursor="second"), + types.ListPromptsResult( + prompts=[ + types.Prompt( + name="changed", + description="New", + arguments=[types.PromptArgument(name="topic", required=True)], + ) + ] + ), + ] + ) + await tool.load_prompts() + + functions_by_name = {function.name: function for function in tool._functions} + assert set(functions_by_name) == {"search", "keep", "changed", "custom"} + assert functions_by_name["search"] is tool_function + assert functions_by_name["custom"] is custom_function + assert functions_by_name["keep"] is kept_prompt + assert functions_by_name["changed"].description == "New" + assert functions_by_name["changed"].parameters()["required"] == ["topic"] + assert tool._progressive_loaded_tool_names == {"keep"} + + +async def test_load_prompts_refresh_forgets_progressive_state_on_normalized_name_collision() -> None: + """A different remote prompt that normalizes to the same local name must not inherit load state.""" + tool = MCPTool(name="docs") # type: ignore[abstract] # ty: ignore[call-non-callable] + tool.session = AsyncMock() + tool.session.list_prompts = AsyncMock( + return_value=types.ListPromptsResult(prompts=[types.Prompt(name="a/b", description="Same")]) + ) + await tool.load_prompts() + old_prompt = tool._functions[0] + assert old_prompt.name == "a-b" + tool._progressive_loaded_tool_names.add("a-b") + + tool.session.list_prompts = AsyncMock( + return_value=types.ListPromptsResult(prompts=[types.Prompt(name="a-b", description="Same")]) + ) + await tool.load_prompts() + + assert len(tool._functions) == 1 + new_prompt = tool._functions[0] + assert new_prompt is not old_prompt + assert new_prompt.name == "a-b" + assert (new_prompt.additional_properties or {})["_mcp_remote_name"] == "a-b" + assert tool._progressive_loaded_tool_names == set() + + +async def test_load_prompts_refresh_keeps_progressive_state_when_metadata_changes() -> None: + """A prompt rebuilt for changed metadata keeps load state while its remote identity is unchanged.""" + tool = MCPTool(name="docs") # type: ignore[abstract] # ty: ignore[call-non-callable] + tool.session = AsyncMock() + tool.session.list_prompts = AsyncMock( + return_value=types.ListPromptsResult(prompts=[types.Prompt(name="summary", description="Old")]) + ) + await tool.load_prompts() + old_prompt = tool._functions[0] + tool._progressive_loaded_tool_names.add("summary") + + tool.session.list_prompts = AsyncMock( + return_value=types.ListPromptsResult( + prompts=[ + types.Prompt( + name="summary", + description="New", + arguments=[types.PromptArgument(name="topic", required=True)], + ) + ] + ) + ) + await tool.load_prompts() + + assert len(tool._functions) == 1 + new_prompt = tool._functions[0] + assert new_prompt is not old_prompt + assert new_prompt.description == "New" + assert "topic" in new_prompt.parameters()["properties"] + assert tool._progressive_loaded_tool_names == {"summary"} + + +async def test_load_prompts_empty_snapshot_clears_prompts() -> None: + """A successful empty prompts/list snapshot clears previously loaded prompts.""" + tool = MCPTool(name="docs") # type: ignore[abstract] # ty: ignore[call-non-callable] + tool.session = AsyncMock() + tool.session.list_prompts = AsyncMock(return_value=types.ListPromptsResult(prompts=[types.Prompt(name="summary")])) + await tool.load_prompts() + custom_function = FunctionTool(name="custom", func=lambda: "custom") + tool._functions.append(custom_function) + + tool.session.list_prompts = AsyncMock(return_value=types.ListPromptsResult(prompts=[])) + await tool.load_prompts() + + assert tool._functions == [custom_function] + + +async def test_load_prompts_failed_pagination_keeps_previous_catalog() -> None: + """A failed page request must not publish a partial prompt snapshot.""" + tool = MCPTool(name="docs") # type: ignore[abstract] # ty: ignore[call-non-callable] + tool.session = AsyncMock() + tool.session.list_prompts = AsyncMock( + return_value=types.ListPromptsResult(prompts=[types.Prompt(name="first"), types.Prompt(name="second")]) + ) + await tool.load_prompts() + original_functions = list(tool._functions) + + tool.session.list_prompts = AsyncMock( + side_effect=[ + types.ListPromptsResult(prompts=[types.Prompt(name="first")], nextCursor="second"), + McpError(types.ErrorData(code=-32603, message="boom")), + ] + ) + with pytest.raises(McpError): + await tool.load_prompts() + + assert tool._functions == original_functions + + def test_mcp_prompt_message_to_ai_content(): """Test conversion from MCP prompt message to AI content.""" mcp_message = types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hello, world!"))