diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index d6a769565..5a45d461a 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -14,7 +14,7 @@ from contextlib import ( ) from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any, Generic, Literal +from typing import TYPE_CHECKING, Any, Generic, Literal, overload import anyio import httpx @@ -45,6 +45,7 @@ import fastmcp.server import fastmcp.settings from fastmcp.exceptions import NotFoundError from fastmcp.prompts import Prompt, PromptManager +from fastmcp.prompts.prompt import FunctionPrompt from fastmcp.resources import Resource, ResourceManager from fastmcp.resources.template import ResourceTemplate from fastmcp.server.auth.auth import OAuthProvider @@ -55,7 +56,7 @@ from fastmcp.server.http import ( create_streamable_http_app, ) from fastmcp.tools import ToolManager -from fastmcp.tools.tool import Tool +from fastmcp.tools.tool import FunctionTool, Tool from fastmcp.utilities.cache import TimedCache from fastmcp.utilities.logging import get_logger from fastmcp.utilities.mcp_config import MCPConfig @@ -510,6 +511,30 @@ class FastMCP(Generic[LifespanResultT]): self._tool_manager.remove_tool(name) self._cache.clear() + @overload + def tool( + self, + name_or_fn: AnyFunction, + *, + name: str | None = None, + description: str | None = None, + tags: set[str] | None = None, + annotations: ToolAnnotations | dict[str, Any] | None = None, + exclude_args: list[str] | None = None, + ) -> FunctionTool: ... + + @overload + def tool( + self, + name_or_fn: str | None = None, + *, + name: str | None = None, + description: str | None = None, + tags: set[str] | None = None, + annotations: ToolAnnotations | dict[str, Any] | None = None, + exclude_args: list[str] | None = None, + ) -> Callable[[AnyFunction], FunctionTool]: ... + def tool( self, name_or_fn: str | AnyFunction | None = None, @@ -519,7 +544,7 @@ class FastMCP(Generic[LifespanResultT]): tags: set[str] | None = None, annotations: ToolAnnotations | dict[str, Any] | None = None, exclude_args: list[str] | None = None, - ) -> Callable[[AnyFunction], AnyFunction] | AnyFunction: + ) -> Callable[[AnyFunction], FunctionTool] | FunctionTool: """Decorator to register a tool. Tools can optionally request a Context object by adding a parameter with the @@ -571,7 +596,7 @@ class FastMCP(Generic[LifespanResultT]): fn = name_or_fn tool_name = name # Use keyword name if provided, otherwise None - # Register the tool immediately and return the function + # Register the tool immediately and return the tool object tool = Tool.from_function( fn, name=tool_name, @@ -582,7 +607,7 @@ class FastMCP(Generic[LifespanResultT]): serializer=self._tool_serializer, ) self.add_tool(tool) - return fn + return tool elif isinstance(name_or_fn, str): # Case 3: @tool("custom_name") - name passed as first argument @@ -674,7 +699,7 @@ class FastMCP(Generic[LifespanResultT]): description: str | None = None, mime_type: str | None = None, tags: set[str] | None = None, - ) -> Callable[[AnyFunction], AnyFunction]: + ) -> Callable[[AnyFunction], Resource | ResourceTemplate]: """Decorator to register a function as a resource. The function will be called when the resource is read to generate its content. @@ -728,7 +753,7 @@ class FastMCP(Generic[LifespanResultT]): "Did you forget to call it? Use @resource('uri') instead of @resource" ) - def decorator(fn: AnyFunction) -> AnyFunction: + def decorator(fn: AnyFunction) -> Resource | ResourceTemplate: from fastmcp.server.context import Context # Check if this should be a template @@ -750,6 +775,7 @@ class FastMCP(Generic[LifespanResultT]): tags=tags, ) self.add_template(template) + return template elif not has_uri_params and not has_func_params: resource = Resource.from_function( fn=fn, @@ -760,14 +786,13 @@ class FastMCP(Generic[LifespanResultT]): tags=tags, ) self.add_resource(resource) + return resource else: raise ValueError( "Invalid resource or template definition due to a " "mismatch between URI parameters and function parameters." ) - return fn - return decorator def add_prompt(self, prompt: Prompt) -> None: @@ -779,6 +804,26 @@ class FastMCP(Generic[LifespanResultT]): self._prompt_manager.add_prompt(prompt) self._cache.clear() + @overload + def prompt( + self, + name_or_fn: AnyFunction, + *, + name: str | None = None, + description: str | None = None, + tags: set[str] | None = None, + ) -> FunctionPrompt: ... + + @overload + def prompt( + self, + name_or_fn: str | None = None, + *, + name: str | None = None, + description: str | None = None, + tags: set[str] | None = None, + ) -> Callable[[AnyFunction], FunctionPrompt]: ... + def prompt( self, name_or_fn: str | AnyFunction | None = None, @@ -786,7 +831,7 @@ class FastMCP(Generic[LifespanResultT]): name: str | None = None, description: str | None = None, tags: set[str] | None = None, - ) -> Callable[[AnyFunction], AnyFunction] | AnyFunction: + ) -> Callable[[AnyFunction], FunctionPrompt] | FunctionPrompt: """Decorator to register a prompt. Prompts can optionally request a Context object by adding a parameter with the @@ -867,7 +912,7 @@ class FastMCP(Generic[LifespanResultT]): ) self.add_prompt(prompt) - return fn + return prompt elif isinstance(name_or_fn, str): # Case 3: @prompt("custom_name") - name passed as first argument diff --git a/tests/server/test_server.py b/tests/server/test_server.py index ab729846b..9892f760f 100644 --- a/tests/server/test_server.py +++ b/tests/server/test_server.py @@ -6,7 +6,7 @@ from pydantic import Field from fastmcp import Client, FastMCP from fastmcp.exceptions import NotFoundError -from fastmcp.prompts.prompt import Prompt +from fastmcp.prompts.prompt import FunctionPrompt, Prompt from fastmcp.resources import Resource, ResourceTemplate from fastmcp.server.server import ( MountedServer, @@ -179,7 +179,6 @@ class TestToolDecorator: def __init__(self, x: int): self.x = x - @mcp.tool def add(self, y: int) -> int: return self.x + y @@ -326,11 +325,11 @@ class TestToolDecorator: result_fn = mcp.tool(standalone_function, name="direct_call_tool") # The function should be returned unchanged - assert result_fn is standalone_function + assert isinstance(result_fn, FunctionTool) # Verify the tool was registered correctly tools = await mcp.get_tools() - assert "direct_call_tool" in tools + assert tools["direct_call_tool"] is result_fn # Verify it can be called result = await mcp._mcp_call_tool("direct_call_tool", {"x": 5, "y": 3}) @@ -857,11 +856,11 @@ class TestPromptDecorator: result_fn = mcp.prompt(standalone_function, name="direct_call_prompt") # The function should be returned unchanged - assert result_fn is standalone_function + assert isinstance(result_fn, FunctionPrompt) # Verify the prompt was registered correctly prompts = await mcp.get_prompts() - assert "direct_call_prompt" in prompts + assert prompts["direct_call_prompt"] is result_fn # Verify it can be called async with Client(mcp) as client: diff --git a/tests/server/test_server_interactions.py b/tests/server/test_server_interactions.py index a15f9b5ff..921902819 100644 --- a/tests/server/test_server_interactions.py +++ b/tests/server/test_server_interactions.py @@ -936,56 +936,6 @@ class TestResourceTemplates: result = await client.read_resource(AnyUrl("resource://test/data")) assert result[0].text == "Data for test" # type: ignore[attr-defined] - async def test_stacked_resource_template_decorators(self): - """Test that resource template decorators can be stacked.""" - mcp = FastMCP() - - @mcp.resource("users://email/{email}") - @mcp.resource("users://name/{name}") - def lookup_user(name: str | None = None, email: str | None = None) -> dict: - if name: - return { - "id": "123", - "name": name, - "email": "dummy@example.com", - "lookup": "name", - } - elif email: - return { - "id": "123", - "name": "Test User", - "email": email, - "lookup": "email", - } - else: - raise ValueError("Either name or email must be provided") - - # Verify both templates are registered - templates_dict = await mcp.get_resource_templates() - templates = list(templates_dict.values()) - assert len(templates) == 2 - template_uris = {t.uri_template for t in templates} - assert "users://email/{email}" in template_uris - assert "users://name/{name}" in template_uris - - # Test lookup by email - async with Client(mcp) as client: - email_result = await client.read_resource( - AnyUrl("users://email/user@example.com") - ) - assert email_result[0].text # type: ignore[attr-defined] - email_data = json.loads(email_result[0].text) # type: ignore[attr-defined] - assert email_data["lookup"] == "email" - assert email_data["email"] == "user@example.com" - - # Test lookup by name - name_result = await client.read_resource(AnyUrl("users://name/John")) - assert name_result[0].text # type: ignore[attr-defined] - name_data = json.loads(name_result[0].text) # type: ignore[attr-defined] - assert name_data["lookup"] == "name" - assert name_data["name"] == "John" - assert name_data["email"] == "dummy@example.com" - async def test_template_decorator_with_tags(self): mcp = FastMCP()