Support enabled/disabled prompts

This commit is contained in:
Jeremiah Lowin 2025-06-10 10:24:40 -04:00
commit 17a71278aa
5 changed files with 103 additions and 4 deletions

View file

@ -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,
)

View file

@ -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."""

View file

@ -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:

View file

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

View file

@ -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()