diff --git a/src/mcp/server/mcpserver/exceptions.py b/src/mcp/server/mcpserver/exceptions.py index dd1b75e82..4320ef180 100644 --- a/src/mcp/server/mcpserver/exceptions.py +++ b/src/mcp/server/mcpserver/exceptions.py @@ -13,6 +13,10 @@ class ResourceError(MCPServerError): """Error in resource operations.""" +class PromptError(MCPServerError): + """Error in prompt operations.""" + + class ToolError(MCPServerError): """Error in tool operations.""" diff --git a/src/mcp/server/mcpserver/prompts/manager.py b/src/mcp/server/mcpserver/prompts/manager.py index 28a7a6e98..2b89d14e6 100644 --- a/src/mcp/server/mcpserver/prompts/manager.py +++ b/src/mcp/server/mcpserver/prompts/manager.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any +from mcp.server.mcpserver.exceptions import PromptError from mcp.server.mcpserver.prompts.base import Message, Prompt from mcp.server.mcpserver.utilities.logging import get_logger @@ -45,6 +46,19 @@ def add_prompt( self._prompts[prompt.name] = prompt return prompt + def remove_prompt(self, name: str) -> None: + """Remove a prompt by name. + + Args: + name: The name of the prompt to remove + + Raises: + PromptError: If the prompt does not exist + """ + if name not in self._prompts: + raise PromptError(f"Unknown prompt: {name}") + del self._prompts[name] + async def render_prompt( self, name: str, diff --git a/src/mcp/server/mcpserver/resources/resource_manager.py b/src/mcp/server/mcpserver/resources/resource_manager.py index 6bf17376d..b5a6149bd 100644 --- a/src/mcp/server/mcpserver/resources/resource_manager.py +++ b/src/mcp/server/mcpserver/resources/resource_manager.py @@ -7,6 +7,7 @@ from pydantic import AnyUrl +from mcp.server.mcpserver.exceptions import ResourceError from mcp.server.mcpserver.resources.base import Resource from mcp.server.mcpserver.resources.templates import ResourceTemplate from mcp.server.mcpserver.utilities.logging import get_logger @@ -108,3 +109,30 @@ def list_templates(self) -> list[ResourceTemplate]: """List all registered templates.""" logger.debug("Listing templates", extra={"count": len(self._templates)}) return list(self._templates.values()) + + def remove_resource(self, uri: AnyUrl | str) -> None: + """Remove a resource by URI. + + Args: + uri: The URI of the resource to remove + + Raises: + ResourceError: If the resource does not exist + """ + uri_str = str(uri) + if uri_str not in self._resources: + raise ResourceError(f"Unknown resource: {uri}") + del self._resources[uri_str] + + def remove_template(self, uri_template: str) -> None: + """Remove a resource template by URI template. + + Args: + uri_template: The URI template of the template to remove + + Raises: + ResourceError: If the template does not exist + """ + if uri_template not in self._templates: + raise ResourceError(f"Unknown template: {uri_template}") + del self._templates[uri_template] diff --git a/src/mcp/server/mcpserver/server.py b/src/mcp/server/mcpserver/server.py index 2a7a58117..80675d3e3 100644 --- a/src/mcp/server/mcpserver/server.py +++ b/src/mcp/server/mcpserver/server.py @@ -500,6 +500,39 @@ def remove_tool(self, name: str) -> None: """ self._tool_manager.remove_tool(name) + def remove_prompt(self, name: str) -> None: + """Remove a prompt from the server by name. + + Args: + name: The name of the prompt to remove + + Raises: + PromptError: If the prompt does not exist + """ + self._prompt_manager.remove_prompt(name) + + def remove_resource(self, uri: str) -> None: + """Remove a resource from the server by URI. + + Args: + uri: The URI of the resource to remove + + Raises: + ResourceError: If the resource does not exist + """ + self._resource_manager.remove_resource(uri) + + def remove_resource_template(self, uri_template: str) -> None: + """Remove a resource template from the server by URI template. + + Args: + uri_template: The URI template of the template to remove + + Raises: + ResourceError: If the template does not exist + """ + self._resource_manager.remove_template(uri_template) + def tool( self, name: str | None = None, diff --git a/tests/server/mcpserver/prompts/test_manager.py b/tests/server/mcpserver/prompts/test_manager.py index 99a03db56..167ab73df 100644 --- a/tests/server/mcpserver/prompts/test_manager.py +++ b/tests/server/mcpserver/prompts/test_manager.py @@ -108,3 +108,70 @@ def fn(name: str) -> str: # pragma: no cover manager.add_prompt(prompt) with pytest.raises(ValueError, match="Missing required arguments"): await manager.render_prompt("fn", None, Context()) + + +class TestRemovePrompt: + """Test PromptManager.remove_prompt() functionality.""" + + def test_remove_existing_prompt(self): + """Test removing an existing prompt.""" + + def fn() -> str: # pragma: no cover + return "Hello, world!" + + manager = PromptManager() + prompt = Prompt.from_function(fn) + manager.add_prompt(prompt) + + # Verify prompt exists + assert manager.get_prompt("fn") is not None + assert len(manager.list_prompts()) == 1 + + # Remove the prompt - should not raise any exception + manager.remove_prompt("fn") + + # Verify prompt is removed + assert manager.get_prompt("fn") is None + assert len(manager.list_prompts()) == 0 + + def test_remove_nonexistent_prompt(self): + """Test removing a non-existent prompt raises error.""" + manager = PromptManager() + + with pytest.raises(Exception, match="Unknown prompt: nonexistent"): + manager.remove_prompt("nonexistent") + + def test_remove_one_prompt_from_multiple(self): + """Test removing one prompt when multiple prompts exist.""" + + def fn1() -> str: # pragma: no cover + return "Hello, world!" + + def fn2() -> str: # pragma: no cover + return "Goodbye, world!" + + def fn3() -> str: # pragma: no cover + return "How are you?" + + manager = PromptManager() + prompt1 = Prompt.from_function(fn1) + prompt2 = Prompt.from_function(fn2) + prompt3 = Prompt.from_function(fn3) + manager.add_prompt(prompt1) + manager.add_prompt(prompt2) + manager.add_prompt(prompt3) + + # Verify all prompts exist + assert len(manager.list_prompts()) == 3 + assert manager.get_prompt("fn1") is not None + assert manager.get_prompt("fn2") is not None + assert manager.get_prompt("fn3") is not None + + # Remove middle prompt + manager.remove_prompt("fn2") + + # Verify only fn2 is removed + assert len(manager.list_prompts()) == 2 + assert manager.get_prompt("fn1") is not None + assert manager.get_prompt("fn2") is None + assert manager.get_prompt("fn3") is not None diff --git a/tests/server/mcpserver/resources/test_resource_manager.py b/tests/server/mcpserver/resources/test_resource_manager.py index 724b57997..bae5a183e 100644 --- a/tests/server/mcpserver/resources/test_resource_manager.py +++ b/tests/server/mcpserver/resources/test_resource_manager.py @@ -5,6 +5,7 @@ from pydantic import AnyUrl from mcp.server.mcpserver import Context +from mcp.server.mcpserver.exceptions import ResourceError from mcp.server.mcpserver.resources import FileResource, FunctionResource, ResourceManager, ResourceTemplate @@ -175,3 +176,165 @@ def get_item(id: str) -> str: # pragma: no cover ) assert template.meta is None + + +class TestRemoveResource: + """Test ResourceManager.remove_resource() functionality.""" + + def test_remove_existing_resource(self, temp_file: Path): + """Test removing an existing resource.""" + manager = ResourceManager() + resource = FileResource( + uri=f"file://{temp_file}", + name="test", + path=temp_file, + ) + manager.add_resource(resource) + + # Verify resource exists + assert len(manager.list_resources()) == 1 + + # Remove the resource - should not raise any exception + manager.remove_resource(f"file://{temp_file}") + + # Verify resource is removed + assert len(manager.list_resources()) == 0 + + def test_remove_nonexistent_resource(self): + """Test removing a non-existent resource raises ResourceError.""" + manager = ResourceManager() + + with pytest.raises(ResourceError, match="Unknown resource: nonexistent"): + manager.remove_resource("nonexistent") + + def test_remove_resource_with_anyurl(self, temp_file: Path): + """Test removing a resource using AnyUrl instead of string.""" + manager = ResourceManager() + resource = FileResource( + uri=f"file://{temp_file}", + name="test", + path=temp_file, + ) + manager.add_resource(resource) + + # Verify resource exists + assert len(manager.list_resources()) == 1 + + # Remove using AnyUrl - should work + manager.remove_resource(AnyUrl(f"file://{temp_file}")) + + # Verify resource is removed + assert len(manager.list_resources()) == 0 + + def test_remove_one_resource_from_multiple(self, temp_file: Path): + """Test removing one resource when multiple resources exist.""" + manager = ResourceManager() + resource1 = FileResource( + uri=f"file://{temp_file}1", + name="test1", + path=temp_file, + ) + resource2 = FileResource( + uri=f"file://{temp_file}2", + name="test2", + path=temp_file, + ) + resource3 = FileResource( + uri=f"file://{temp_file}3", + name="test3", + path=temp_file, + ) + manager.add_resource(resource1) + manager.add_resource(resource2) + manager.add_resource(resource3) + + # Verify all resources exist + assert len(manager.list_resources()) == 3 + + # Remove middle resource + manager.remove_resource(f"file://{temp_file}2") + + # Verify only resource2 is removed + assert len(manager.list_resources()) == 2 + assert manager.list_resources() == [resource1, resource3] + + @pytest.mark.anyio + async def test_call_removed_resource_raises_error(self, temp_file: Path): + """Test that calling a removed resource raises ValueError.""" + manager = ResourceManager() + resource = FileResource( + uri=f"file://{temp_file}", + name="test", + path=temp_file, + ) + manager.add_resource(resource) + + # Verify resource works before removal + result = await manager.get_resource(resource.uri, Context()) + assert result == resource + + # Remove the resource + manager.remove_resource(f"file://{temp_file}") + + # Verify getting removed resource raises error + with pytest.raises(ValueError, match="Unknown resource"): + await manager.get_resource(resource.uri, Context()) + + +class TestRemoveTemplate: + """Test ResourceManager.remove_template() functionality.""" + + def test_remove_existing_template(self): + """Test removing an existing template.""" + manager = ResourceManager() + + def greet(name: str) -> str: + return f"Hello, {name}!" + + manager.add_template( + fn=greet, + uri_template="greet://{name}", + ) + + # Verify template exists + assert len(manager.list_templates()) == 1 + + # Remove the template + manager.remove_template("greet://{name}") + + # Verify template is removed + assert len(manager.list_templates()) == 0 + + def test_remove_nonexistent_template(self): + """Test removing a non-existent template raises ResourceError.""" + manager = ResourceManager() + + with pytest.raises(ResourceError, match="Unknown template: nonexistent"): + manager.remove_template("nonexistent") + + def test_remove_one_template_from_multiple(self): + """Test removing one template when multiple templates exist.""" + manager = ResourceManager() + + def greet(name: str) -> str: + return f"Hello, {name}!" + + def farewell(name: str) -> str: + return f"Goodbye, {name}!" + + def ask(question: str) -> str: + return f"What is {question}?" + + template1 = manager.add_template(fn=greet, uri_template="greet://{name}") + manager.add_template(fn=farewell, uri_template="farewell://{name}") + template3 = manager.add_template(fn=ask, uri_template="ask://{question}") + + # Verify all templates exist + assert len(manager.list_templates()) == 3 + + # Remove middle template + manager.remove_template("farewell://{name}") + + # Verify only farewell template is removed + assert len(manager.list_templates()) == 2 + assert manager.list_templates() == [template1, template3] diff --git a/tests/server/mcpserver/test_server.py b/tests/server/mcpserver/test_server.py index 3ef06d038..a034ed485 100644 --- a/tests/server/mcpserver/test_server.py +++ b/tests/server/mcpserver/test_server.py @@ -1481,3 +1481,141 @@ async def test_report_progress_passes_related_request_id(): message="halfway", related_request_id="req-abc-123", ) + + +class TestRemovePrompt: + """Test remove_prompt functionality in MCPServer.""" + + async def test_remove_prompt(self): + """Test removing a prompt from the server.""" + mcp = MCPServer() + + @mcp.prompt() + def fn() -> str: + return "Hello, world!" + + # Verify prompt exists + assert len(mcp._prompt_manager.list_prompts()) == 1 + + # Remove the prompt + mcp.remove_prompt("fn") + + # Verify prompt is removed + assert len(mcp._prompt_manager.list_prompts()) == 0 + + async def test_remove_nonexistent_prompt(self): + """Test that removing a non-existent prompt raises PromptError.""" + from mcp.server.mcpserver.exceptions import PromptError + + mcp = MCPServer() + + with pytest.raises(PromptError, match="Unknown prompt: nonexistent"): + mcp.remove_prompt("nonexistent") + + async def test_remove_prompt_and_list(self): + """Test that a removed prompt doesn't appear in list_prompts.""" + mcp = MCPServer() + + @mcp.prompt() + def prompt1() -> str: + return "Prompt 1" + + @mcp.prompt() + def prompt2() -> str: + return "Prompt 2" + + # Verify both prompts exist + async with Client(mcp) as client: + prompts = await client.list_prompts() + assert len(prompts.prompts) == 2 + + # Remove one prompt + mcp.remove_prompt("prompt1") + + # Verify only one prompt remains + async with Client(mcp) as client: + prompts = await client.list_prompts() + assert len(prompts.prompts) == 1 + assert prompts.prompts[0].name == "prompt2" + + +class TestRemoveResource: + """Test remove_resource functionality in MCPServer.""" + + async def test_remove_resource(self): + """Test removing a resource from the server.""" + mcp = MCPServer() + + @mcp.resource("resource://test") + def get_resource() -> str: + return "Hello, world!" + + # Verify resource exists + assert len(mcp._resource_manager.list_resources()) == 1 + + # Remove the resource + mcp.remove_resource("resource://test") + + # Verify resource is removed + assert len(mcp._resource_manager.list_resources()) == 0 + + async def test_remove_nonexistent_resource(self): + """Test that removing a non-existent resource raises ResourceError.""" + from mcp.server.mcpserver.exceptions import ResourceError + + mcp = MCPServer() + + with pytest.raises(ResourceError, match="Unknown resource: nonexistent"): + mcp.remove_resource("nonexistent") + + async def test_remove_resource_and_list(self): + """Test that a removed resource doesn't appear in list_resources.""" + mcp = MCPServer() + + @mcp.resource("resource://test1") + def resource1() -> str: + return "Resource 1" + + @mcp.resource("resource://test2") + def resource2() -> str: + return "Resource 2" + + # Verify both resources exist + async with Client(mcp) as client: + resources = await client.list_resources() + assert len(resources.resources) == 2 + + # Remove one resource + mcp.remove_resource("resource://test1") + + # Verify only one resource remains + async with Client(mcp) as client: + resources = await client.list_resources() + assert len(resources.resources) == 1 + assert resources.resources[0].uri == "resource://test2" + + async def test_remove_resource_template(self): + """Test removing a resource template from the server.""" + mcp = MCPServer() + + @mcp.resource("resource://{name}") + def get_resource(name: str) -> str: + return f"Hello, {name}!" + + # Verify template exists + assert len(mcp._resource_manager.list_templates()) == 1 + + # Remove the template + mcp.remove_resource_template("resource://{name}") + + # Verify template is removed + assert len(mcp._resource_manager.list_templates()) == 0 + + async def test_remove_nonexistent_template(self): + """Test that removing a non-existent template raises ResourceError.""" + from mcp.server.mcpserver.exceptions import ResourceError + + mcp = MCPServer() + + with pytest.raises(ResourceError, match="Unknown template: nonexistent"): + mcp.remove_resource_template("nonexistent")