From 17a71278aaa30e017b44d528b42fc7b9ea84b33e Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 10 Jun 2025 10:24:40 -0400 Subject: [PATCH] Support enabled/disabled prompts --- src/fastmcp/prompts/prompt.py | 5 +- src/fastmcp/prompts/prompt_manager.py | 6 +- src/fastmcp/server/server.py | 5 ++ tests/server/test_proxy.py | 4 +- tests/server/test_server_interactions.py | 87 ++++++++++++++++++++++++ 5 files changed, 103 insertions(+), 4 deletions(-) diff --git a/src/fastmcp/prompts/prompt.py b/src/fastmcp/prompts/prompt.py index 63b276745..a6c93cd52 100644 --- a/src/fastmcp/prompts/prompt.py +++ b/src/fastmcp/prompts/prompt.py @@ -96,6 +96,7 @@ class Prompt(FastMCPComponent, ABC): name: str | None = None, description: str | None = None, tags: set[str] | None = None, + enabled: bool | None = None, ) -> FunctionPrompt: """Create a Prompt from a function. @@ -106,7 +107,7 @@ class Prompt(FastMCPComponent, ABC): - A sequence of any of the above """ return FunctionPrompt.from_function( - fn=fn, name=name, description=description, tags=tags + fn=fn, name=name, description=description, tags=tags, enabled=enabled ) @abstractmethod @@ -130,6 +131,7 @@ class FunctionPrompt(Prompt): name: str | None = None, description: str | None = None, tags: set[str] | None = None, + enabled: bool | None = None, ) -> FunctionPrompt: """Create a Prompt from a function. @@ -195,6 +197,7 @@ class FunctionPrompt(Prompt): description=description, arguments=arguments, tags=tags or set(), + enabled=enabled if enabled is not None else True, fn=fn, ) diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index fea5ebcd4..4431eadab 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -40,9 +40,11 @@ class PromptManager: self.duplicate_behavior = duplicate_behavior - def get_prompt(self, key: str) -> Prompt | None: + def get_prompt(self, key: str) -> Prompt: """Get prompt by key.""" - return self._prompts.get(key) + if key in self._prompts: + return self._prompts[key] + raise NotFoundError(f"Unknown prompt: {key}") def get_prompts(self) -> dict[str, Prompt]: """Get all registered prompts, indexed by registered key.""" diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index d7922f84b..0198bc6c6 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -940,6 +940,7 @@ class FastMCP(Generic[LifespanResultT]): name: str | None = None, description: str | None = None, tags: set[str] | None = None, + enabled: bool | None = None, ) -> FunctionPrompt: ... @overload @@ -950,6 +951,7 @@ class FastMCP(Generic[LifespanResultT]): name: str | None = None, description: str | None = None, tags: set[str] | None = None, + enabled: bool | None = None, ) -> Callable[[AnyFunction], FunctionPrompt]: ... def prompt( @@ -959,6 +961,7 @@ class FastMCP(Generic[LifespanResultT]): name: str | None = None, description: str | None = None, tags: set[str] | None = None, + enabled: bool | None = None, ) -> Callable[[AnyFunction], FunctionPrompt] | FunctionPrompt: """Decorator to register a prompt. @@ -1050,6 +1053,7 @@ class FastMCP(Generic[LifespanResultT]): name=prompt_name, description=description, tags=tags, + enabled=enabled, ) self.add_prompt(prompt) @@ -1077,6 +1081,7 @@ class FastMCP(Generic[LifespanResultT]): name=prompt_name, description=description, tags=tags, + enabled=enabled, ) async def run_stdio_async(self) -> None: diff --git a/tests/server/test_proxy.py b/tests/server/test_proxy.py index 83b5618cd..164f33bf8 100644 --- a/tests/server/test_proxy.py +++ b/tests/server/test_proxy.py @@ -178,7 +178,9 @@ class TestResources: assert json.loads(result[0].text) == USERS # type: ignore[attr-defined] async def test_read_resource_returns_none_if_not_found(self, proxy_server): - with pytest.raises(McpError, match="Unknown resource: resource://nonexistent"): + with pytest.raises( + McpError, match="Unknown resource: 'resource://nonexistent'" + ): async with Client(proxy_server) as client: await client.read_resource("resource://nonexistent") diff --git a/tests/server/test_server_interactions.py b/tests/server/test_server_interactions.py index a1a8de58f..124474703 100644 --- a/tests/server/test_server_interactions.py +++ b/tests/server/test_server_interactions.py @@ -1220,6 +1220,93 @@ class TestPrompts: assert prompt.tags == {"example", "test-tag"} +class TestPromptEnabled: + async def test_toggle_enabled(self): + mcp = FastMCP() + + @mcp.prompt + def sample_prompt() -> str: + return "Hello, world!" + + assert sample_prompt.enabled + + prompt = await mcp.get_prompt("sample_prompt") + assert prompt.enabled + + prompt.disable() + + assert not prompt.enabled + assert not sample_prompt.enabled + + prompt.enable() + assert prompt.enabled + assert sample_prompt.enabled + + async def test_prompt_disabled_in_decorator(self): + mcp = FastMCP() + + @mcp.prompt(enabled=False) + def sample_prompt() -> str: + return "Hello, world!" + + async with Client(mcp) as client: + prompts = await client.list_prompts() + assert len(prompts) == 0 + + async def test_prompt_toggle_enabled(self): + mcp = FastMCP() + + @mcp.prompt(enabled=False) + def sample_prompt() -> str: + return "Hello, world!" + + sample_prompt.enable() + + async with Client(mcp) as client: + prompts = await client.list_prompts() + assert len(prompts) == 1 + + async def test_prompt_toggle_disabled(self): + mcp = FastMCP() + + @mcp.prompt + def sample_prompt() -> str: + return "Hello, world!" + + sample_prompt.disable() + + async with Client(mcp) as client: + prompts = await client.list_prompts() + assert len(prompts) == 0 + + async def test_get_prompt_and_disable(self): + mcp = FastMCP() + + @mcp.prompt + def sample_prompt() -> str: + return "Hello, world!" + + prompt = await mcp.get_prompt("sample_prompt") + assert prompt.enabled + + sample_prompt.disable() + + async with Client(mcp) as client: + result = await client.list_prompts() + assert len(result) == 0 + + async def test_cant_get_disabled_prompt(self): + mcp = FastMCP() + + @mcp.prompt(enabled=False) + def sample_prompt() -> str: + return "Hello, world!" + + with pytest.raises(McpError, match="Unknown prompt"): + async with Client(mcp) as client: + await client.get_prompt("sample_prompt") + + class TestPromptContext: async def test_prompt_context(self): mcp = FastMCP()