Python: Replace MCP prompt catalog on refresh

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.

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!"))