diff --git a/src/fastmcp/__init__.py b/src/fastmcp/__init__.py index ddcb3c456..083fc21d9 100644 --- a/src/fastmcp/__init__.py +++ b/src/fastmcp/__init__.py @@ -1,12 +1,11 @@ """FastMCP - An ergonomic MCP interface.""" from importlib.metadata import version -from fastmcp.server import FastMCP, Context +import fastmcp.settings + +from fastmcp.server.server import FastMCP +from fastmcp.server.context import Context from . import clients __version__ = version("fastmcp") -__all__ = [ - "FastMCP", - "Context", - "clients", -] +__all__ = ["FastMCP", "Context", "clients"] diff --git a/src/fastmcp/exceptions.py b/src/fastmcp/exceptions.py new file mode 100644 index 000000000..fb5bda106 --- /dev/null +++ b/src/fastmcp/exceptions.py @@ -0,0 +1,21 @@ +"""Custom exceptions for FastMCP.""" + + +class FastMCPError(Exception): + """Base error for FastMCP.""" + + +class ValidationError(FastMCPError): + """Error in validating parameters or return values.""" + + +class ResourceError(FastMCPError): + """Error in resource operations.""" + + +class ToolError(FastMCPError): + """Error in tool operations.""" + + +class InvalidSignature(Exception): + """Invalid signature for use with FastMCP.""" diff --git a/src/fastmcp/prompts/__init__.py b/src/fastmcp/prompts/__init__.py index 0c53974a5..bacb4c37d 100644 --- a/src/fastmcp/prompts/__init__.py +++ b/src/fastmcp/prompts/__init__.py @@ -1,3 +1,4 @@ +from .base import Prompt from .prompt_manager import PromptManager -__all__ = ["PromptManager"] +__all__ = ["Prompt", "PromptManager"] diff --git a/src/fastmcp/prompts/base.py b/src/fastmcp/prompts/base.py new file mode 100644 index 000000000..98d141786 --- /dev/null +++ b/src/fastmcp/prompts/base.py @@ -0,0 +1,166 @@ +"""Base classes for FastMCP prompts.""" + +import inspect +import json +from collections.abc import Awaitable, Callable, Sequence +from typing import Any, Literal + +import pydantic_core +from mcp.types import EmbeddedResource, ImageContent, TextContent +from pydantic import BaseModel, Field, TypeAdapter, validate_call + +CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource + + +class Message(BaseModel): + """Base class for all prompt messages.""" + + role: Literal["user", "assistant"] + content: CONTENT_TYPES + + def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any): + if isinstance(content, str): + content = TextContent(type="text", text=content) + super().__init__(content=content, **kwargs) + + +class UserMessage(Message): + """A message from the user.""" + + role: Literal["user", "assistant"] = "user" + + def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any): + super().__init__(content=content, **kwargs) + + +class AssistantMessage(Message): + """A message from the assistant.""" + + role: Literal["user", "assistant"] = "assistant" + + def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any): + super().__init__(content=content, **kwargs) + + +message_validator = TypeAdapter[UserMessage | AssistantMessage]( + UserMessage | AssistantMessage +) + +SyncPromptResult = ( + str | Message | dict[str, Any] | Sequence[str | Message | dict[str, Any]] +) +PromptResult = SyncPromptResult | Awaitable[SyncPromptResult] + + +class PromptArgument(BaseModel): + """An argument that can be passed to a prompt.""" + + name: str = Field(description="Name of the argument") + description: str | None = Field( + None, description="Description of what the argument does" + ) + required: bool = Field( + default=False, description="Whether the argument is required" + ) + + +class Prompt(BaseModel): + """A prompt template that can be rendered with parameters.""" + + name: str = Field(description="Name of the prompt") + description: str | None = Field( + None, description="Description of what the prompt does" + ) + arguments: list[PromptArgument] | None = Field( + None, description="Arguments that can be passed to the prompt" + ) + fn: Callable[..., PromptResult | Awaitable[PromptResult]] = Field(exclude=True) + + @classmethod + def from_function( + cls, + fn: Callable[..., PromptResult | Awaitable[PromptResult]], + name: str | None = None, + description: str | None = None, + ) -> "Prompt": + """Create a Prompt from a function. + + The function can return: + - A string (converted to a message) + - A Message object + - A dict (converted to a message) + - A sequence of any of the above + """ + func_name = name or fn.__name__ + + if func_name == "": + raise ValueError("You must provide a name for lambda functions") + + # Get schema from TypeAdapter - will fail if function isn't properly typed + parameters = TypeAdapter(fn).json_schema() + + # Convert parameters to PromptArguments + arguments: list[PromptArgument] = [] + if "properties" in parameters: + for param_name, param in parameters["properties"].items(): + required = param_name in parameters.get("required", []) + arguments.append( + PromptArgument( + name=param_name, + description=param.get("description"), + required=required, + ) + ) + + # ensure the arguments are properly cast + fn = validate_call(fn) + + return cls( + name=func_name, + description=description or fn.__doc__ or "", + arguments=arguments, + fn=fn, + ) + + async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]: + """Render the prompt with arguments.""" + # Validate required arguments + if self.arguments: + required = {arg.name for arg in self.arguments if arg.required} + provided = set(arguments or {}) + missing = required - provided + if missing: + raise ValueError(f"Missing required arguments: {missing}") + + try: + # Call function and check if result is a coroutine + result = self.fn(**(arguments or {})) + if inspect.iscoroutine(result): + result = await result + + # Validate messages + if not isinstance(result, list | tuple): + result = [result] + + # Convert result to messages + messages: list[Message] = [] + for msg in result: # type: ignore[reportUnknownVariableType] + try: + if isinstance(msg, Message): + messages.append(msg) + elif isinstance(msg, dict): + messages.append(message_validator.validate_python(msg)) + elif isinstance(msg, str): + content = TextContent(type="text", text=msg) + messages.append(UserMessage(content=content)) + else: + content = json.dumps(pydantic_core.to_jsonable_python(msg)) + messages.append(Message(role="user", content=content)) + except Exception: + raise ValueError( + f"Could not convert prompt result to message: {msg}" + ) + + return messages + except Exception as e: + raise ValueError(f"Error rendering prompt {self.name}: {e}") diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index 99317f259..e15da430e 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -1,15 +1,53 @@ -import logging +"""Prompt management functionality.""" -from mcp.server.fastmcp.prompts import PromptManager as BasePromptManager +from typing import Any -logger = logging.getLogger(__name__) +from fastmcp.prompts.base import Message, Prompt +from fastmcp.utilities.logging import get_logger + +logger = get_logger(__name__) -class PromptManager(BasePromptManager): - """ - Extended PromptManager that supports importing prompts from other managers. - Adds ability to import prompts from other managers with prefixed names. - """ +class PromptManager: + """Manages FastMCP prompts.""" + + def __init__(self, warn_on_duplicate_prompts: bool = True): + self._prompts: dict[str, Prompt] = {} + self.warn_on_duplicate_prompts = warn_on_duplicate_prompts + + def get_prompt(self, name: str) -> Prompt | None: + """Get prompt by name.""" + return self._prompts.get(name) + + def list_prompts(self) -> list[Prompt]: + """List all registered prompts.""" + return list(self._prompts.values()) + + def add_prompt( + self, + prompt: Prompt, + ) -> Prompt: + """Add a prompt to the manager.""" + + # Check for duplicates + existing = self._prompts.get(prompt.name) + if existing: + if self.warn_on_duplicate_prompts: + logger.warning(f"Prompt already exists: {prompt.name}") + return existing + + self._prompts[prompt.name] = prompt + return prompt + + async def render_prompt( + self, name: str, arguments: dict[str, Any] | None = None + ) -> list[Message]: + """Render a prompt by name with arguments.""" + prompt = self.get_prompt(name) + if not prompt: + raise ValueError(f"Unknown prompt: {name}") + + return await prompt.render(arguments) def import_prompts( self, manager: "PromptManager", prefix: str | None = None diff --git a/src/fastmcp/resources/__init__.py b/src/fastmcp/resources/__init__.py index 3a13af69d..b5805fb34 100644 --- a/src/fastmcp/resources/__init__.py +++ b/src/fastmcp/resources/__init__.py @@ -1,3 +1,23 @@ +from .base import Resource from .resource_manager import ResourceManager +from .templates import ResourceTemplate +from .types import ( + BinaryResource, + DirectoryResource, + FileResource, + FunctionResource, + HttpResource, + TextResource, +) -__all__ = ["ResourceManager"] +__all__ = [ + "Resource", + "TextResource", + "BinaryResource", + "FunctionResource", + "FileResource", + "HttpResource", + "DirectoryResource", + "ResourceTemplate", + "ResourceManager", +] diff --git a/src/fastmcp/resources/base.py b/src/fastmcp/resources/base.py new file mode 100644 index 000000000..b2050e7f8 --- /dev/null +++ b/src/fastmcp/resources/base.py @@ -0,0 +1,48 @@ +"""Base classes and interfaces for FastMCP resources.""" + +import abc +from typing import Annotated + +from pydantic import ( + AnyUrl, + BaseModel, + ConfigDict, + Field, + UrlConstraints, + ValidationInfo, + field_validator, +) + + +class Resource(BaseModel, abc.ABC): + """Base class for all resources.""" + + model_config = ConfigDict(validate_default=True) + + uri: Annotated[AnyUrl, UrlConstraints(host_required=False)] = Field( + default=..., description="URI of the resource" + ) + name: str | None = Field(description="Name of the resource", default=None) + description: str | None = Field( + description="Description of the resource", default=None + ) + mime_type: str = Field( + default="text/plain", + description="MIME type of the resource content", + pattern=r"^[a-zA-Z0-9]+/[a-zA-Z0-9\-+.]+$", + ) + + @field_validator("name", mode="before") + @classmethod + def set_default_name(cls, name: str | None, info: ValidationInfo) -> str: + """Set default name from URI if not provided.""" + if name: + return name + if uri := info.data.get("uri"): + return str(uri) + raise ValueError("Either name or uri must be provided") + + @abc.abstractmethod + async def read(self) -> str | bytes: + """Read the resource content.""" + pass diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 6fd51b7ea..c3bc06f19 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -1,14 +1,98 @@ -import logging +"""Resource manager functionality.""" -from mcp.server.fastmcp.resources import ( - ResourceManager as BaseResourceManager, -) +from collections.abc import Callable +from typing import Any -logger = logging.getLogger(__name__) +from pydantic import AnyUrl + +from fastmcp.resources.base import Resource +from fastmcp.resources.templates import ResourceTemplate +from fastmcp.utilities.logging import get_logger + +logger = get_logger(__name__) -class ResourceManager(BaseResourceManager): - """ResourceManager that adds methods to import resources from other managers.""" +class ResourceManager: + """Manages FastMCP resources.""" + + def __init__(self, warn_on_duplicate_resources: bool = True): + self._resources: dict[str, Resource] = {} + self._templates: dict[str, ResourceTemplate] = {} + self.warn_on_duplicate_resources = warn_on_duplicate_resources + + def add_resource(self, resource: Resource) -> Resource: + """Add a resource to the manager. + + Args: + resource: A Resource instance to add + + Returns: + The added resource. If a resource with the same URI already exists, + returns the existing resource. + """ + logger.debug( + "Adding resource", + extra={ + "uri": resource.uri, + "type": type(resource).__name__, + "resource_name": resource.name, + }, + ) + existing = self._resources.get(str(resource.uri)) + if existing: + if self.warn_on_duplicate_resources: + logger.warning(f"Resource already exists: {resource.uri}") + return existing + self._resources[str(resource.uri)] = resource + return resource + + def add_template( + self, + fn: Callable[..., Any], + uri_template: str, + name: str | None = None, + description: str | None = None, + mime_type: str | None = None, + ) -> ResourceTemplate: + """Add a template from a function.""" + template = ResourceTemplate.from_function( + fn, + uri_template=uri_template, + name=name, + description=description, + mime_type=mime_type, + ) + self._templates[template.uri_template] = template + return template + + async def get_resource(self, uri: AnyUrl | str) -> Resource | None: + """Get resource by URI, checking concrete resources first, then templates.""" + uri_str = str(uri) + logger.debug("Getting resource", extra={"uri": uri_str}) + + # First check concrete resources + if resource := self._resources.get(uri_str): + return resource + + # Then check templates + for template in self._templates.values(): + if params := template.matches(uri_str): + try: + return await template.create_resource(uri_str, params) + except Exception as e: + raise ValueError(f"Error creating resource from template: {e}") + + raise ValueError(f"Unknown resource: {uri}") + + def list_resources(self) -> list[Resource]: + """List all registered resources.""" + logger.debug("Listing resources", extra={"count": len(self._resources)}) + return list(self._resources.values()) + + 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 import_resources( self, manager: "ResourceManager", prefix: str | None = None diff --git a/src/fastmcp/resources/templates.py b/src/fastmcp/resources/templates.py new file mode 100644 index 000000000..ac1316e73 --- /dev/null +++ b/src/fastmcp/resources/templates.py @@ -0,0 +1,85 @@ +"""Resource template functionality.""" + +from __future__ import annotations + +import inspect +import re +from collections.abc import Callable +from typing import Any + +from pydantic import BaseModel, Field, TypeAdapter, validate_call + +from fastmcp.resources.types import FunctionResource, Resource + + +class ResourceTemplate(BaseModel): + """A template for dynamically creating resources.""" + + uri_template: str = Field( + description="URI template with parameters (e.g. weather://{city}/current)" + ) + name: str = Field(description="Name of the resource") + description: str | None = Field(description="Description of what the resource does") + mime_type: str = Field( + default="text/plain", description="MIME type of the resource content" + ) + fn: Callable[..., Any] = Field(exclude=True) + parameters: dict[str, Any] = Field( + description="JSON schema for function parameters" + ) + + @classmethod + def from_function( + cls, + fn: Callable[..., Any], + uri_template: str, + name: str | None = None, + description: str | None = None, + mime_type: str | None = None, + ) -> ResourceTemplate: + """Create a template from a function.""" + func_name = name or fn.__name__ + if func_name == "": + raise ValueError("You must provide a name for lambda functions") + + # Get schema from TypeAdapter - will fail if function isn't properly typed + parameters = TypeAdapter(fn).json_schema() + + # ensure the arguments are properly cast + fn = validate_call(fn) + + return cls( + uri_template=uri_template, + name=func_name, + description=description or fn.__doc__ or "", + mime_type=mime_type or "text/plain", + fn=fn, + parameters=parameters, + ) + + def matches(self, uri: str) -> dict[str, Any] | None: + """Check if URI matches template and extract parameters.""" + # Convert template to regex pattern + pattern = self.uri_template.replace("{", "(?P<").replace("}", ">[^/]+)") + match = re.match(f"^{pattern}$", uri) + if match: + return match.groupdict() + return None + + async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: + """Create a resource from the template with the given parameters.""" + try: + # Call function and check if result is a coroutine + result = self.fn(**params) + if inspect.iscoroutine(result): + result = await result + + return FunctionResource( + uri=uri, # type: ignore + name=self.name, + description=self.description, + mime_type=self.mime_type, + fn=lambda: result, # Capture result in closure + ) + except Exception as e: + raise ValueError(f"Error creating resource from template: {e}") diff --git a/src/fastmcp/resources/types.py b/src/fastmcp/resources/types.py new file mode 100644 index 000000000..89a142395 --- /dev/null +++ b/src/fastmcp/resources/types.py @@ -0,0 +1,185 @@ +"""Concrete resource implementations.""" + +import inspect +import json +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import anyio +import anyio.to_thread +import httpx +import pydantic.json +import pydantic_core +from pydantic import Field, ValidationInfo + +from fastmcp.resources.base import Resource + + +class TextResource(Resource): + """A resource that reads from a string.""" + + text: str = Field(description="Text content of the resource") + + async def read(self) -> str: + """Read the text content.""" + return self.text + + +class BinaryResource(Resource): + """A resource that reads from bytes.""" + + data: bytes = Field(description="Binary content of the resource") + + async def read(self) -> bytes: + """Read the binary content.""" + return self.data + + +class FunctionResource(Resource): + """A resource that defers data loading by wrapping a function. + + The function is only called when the resource is read, allowing for lazy loading + of potentially expensive data. This is particularly useful when listing resources, + as the function won't be called until the resource is actually accessed. + + The function can return: + - str for text content (default) + - bytes for binary content + - other types will be converted to JSON + """ + + fn: Callable[[], Any] = Field(exclude=True) + + async def read(self) -> str | bytes: + """Read the resource by calling the wrapped function.""" + try: + result = ( + await self.fn() if inspect.iscoroutinefunction(self.fn) else self.fn() + ) + if isinstance(result, Resource): + return await result.read() + if isinstance(result, bytes): + return result + if isinstance(result, str): + return result + try: + return json.dumps(pydantic_core.to_jsonable_python(result)) + except (TypeError, pydantic_core.PydanticSerializationError): + # If JSON serialization fails, try str() + return str(result) + except Exception as e: + raise ValueError(f"Error reading resource {self.uri}: {e}") + + +class FileResource(Resource): + """A resource that reads from a file. + + Set is_binary=True to read file as binary data instead of text. + """ + + path: Path = Field(description="Path to the file") + is_binary: bool = Field( + default=False, + description="Whether to read the file as binary data", + ) + mime_type: str = Field( + default="text/plain", + description="MIME type of the resource content", + ) + + @pydantic.field_validator("path") + @classmethod + def validate_absolute_path(cls, path: Path) -> Path: + """Ensure path is absolute.""" + if not path.is_absolute(): + raise ValueError("Path must be absolute") + return path + + @pydantic.field_validator("is_binary") + @classmethod + def set_binary_from_mime_type(cls, is_binary: bool, info: ValidationInfo) -> bool: + """Set is_binary based on mime_type if not explicitly set.""" + if is_binary: + return True + mime_type = info.data.get("mime_type", "text/plain") + return not mime_type.startswith("text/") + + async def read(self) -> str | bytes: + """Read the file content.""" + try: + if self.is_binary: + return await anyio.to_thread.run_sync(self.path.read_bytes) + return await anyio.to_thread.run_sync(self.path.read_text) + except Exception as e: + raise ValueError(f"Error reading file {self.path}: {e}") + + +class HttpResource(Resource): + """A resource that reads from an HTTP endpoint.""" + + url: str = Field(description="URL to fetch content from") + mime_type: str = Field( + default="application/json", description="MIME type of the resource content" + ) + + async def read(self) -> str | bytes: + """Read the HTTP content.""" + async with httpx.AsyncClient() as client: + response = await client.get(self.url) + response.raise_for_status() + return response.text + + +class DirectoryResource(Resource): + """A resource that lists files in a directory.""" + + path: Path = Field(description="Path to the directory") + recursive: bool = Field( + default=False, description="Whether to list files recursively" + ) + pattern: str | None = Field( + default=None, description="Optional glob pattern to filter files" + ) + mime_type: str = Field( + default="application/json", description="MIME type of the resource content" + ) + + @pydantic.field_validator("path") + @classmethod + def validate_absolute_path(cls, path: Path) -> Path: + """Ensure path is absolute.""" + if not path.is_absolute(): + raise ValueError("Path must be absolute") + return path + + def list_files(self) -> list[Path]: + """List files in the directory.""" + if not self.path.exists(): + raise FileNotFoundError(f"Directory not found: {self.path}") + if not self.path.is_dir(): + raise NotADirectoryError(f"Not a directory: {self.path}") + + try: + if self.pattern: + return ( + list(self.path.glob(self.pattern)) + if not self.recursive + else list(self.path.rglob(self.pattern)) + ) + return ( + list(self.path.glob("*")) + if not self.recursive + else list(self.path.rglob("*")) + ) + except Exception as e: + raise ValueError(f"Error listing directory {self.path}: {e}") + + async def read(self) -> str: # Always returns JSON string + """Read the directory listing.""" + try: + files = await anyio.to_thread.run_sync(self.list_files) + file_list = [str(f.relative_to(self.path)) for f in files if f.is_file()] + return json.dumps({"files": file_list}, indent=2) + except Exception as e: + raise ValueError(f"Error reading directory {self.path}: {e}") diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index f0a1355af..7c099cc2a 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -1,7 +1,9 @@ -from typing import Any +from __future__ import annotations as _annotations -import mcp.server.fastmcp -from mcp.server.fastmcp.utilities.logging import get_logger +from collections.abc import Iterable +from typing import Any, Generic, Literal + +from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.server.session import ServerSessionT from mcp.shared.context import LifespanContextT, RequestContext from mcp.types import ( @@ -10,19 +12,168 @@ from mcp.types import ( SamplingMessage, TextContent, ) +from pydantic import BaseModel +from pydantic.networks import AnyUrl + +from fastmcp.server.server import FastMCP +from fastmcp.utilities.logging import get_logger logger = get_logger(__name__) -class Context(mcp.server.fastmcp.Context[ServerSessionT, LifespanContextT]): +class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]): + """Context object providing access to MCP capabilities. + + This provides a cleaner interface to MCP's RequestContext functionality. + It gets injected into tool and resource functions that request it via type hints. + + To use context in a tool function, add a parameter with the Context type annotation: + + ```python + @server.tool() + def my_tool(x: int, ctx: Context) -> str: + # Log messages to the client + ctx.info(f"Processing {x}") + ctx.debug("Debug info") + ctx.warning("Warning message") + ctx.error("Error message") + + # Report progress + ctx.report_progress(50, 100) + + # Access resources + data = ctx.read_resource("resource://data") + + # Get request info + request_id = ctx.request_id + client_id = ctx.client_id + + return str(x) + ``` + + The context parameter name can be anything as long as it's annotated with Context. + The context is optional - tools that don't need it can omit the parameter. + """ + + _request_context: RequestContext[ServerSessionT, LifespanContextT] | None + _fastmcp: FastMCP | None + def __init__( self, *, request_context: RequestContext[ServerSessionT, LifespanContextT] | None = None, - fastmcp: mcp.server.fastmcp.FastMCP | None = None, + fastmcp: FastMCP | None = None, **kwargs: Any, ): - super().__init__(request_context=request_context, fastmcp=fastmcp, **kwargs) + super().__init__(**kwargs) + self._request_context = request_context + self._fastmcp = fastmcp + + @property + def fastmcp(self) -> FastMCP: + """Access to the FastMCP server.""" + if self._fastmcp is None: + raise ValueError("Context is not available outside of a request") + return self._fastmcp + + @property + def request_context(self) -> RequestContext[ServerSessionT, LifespanContextT]: + """Access to the underlying request context.""" + if self._request_context is None: + raise ValueError("Context is not available outside of a request") + return self._request_context + + async def report_progress( + self, progress: float, total: float | None = None + ) -> None: + """Report progress for the current operation. + + Args: + progress: Current progress value e.g. 24 + total: Optional total value e.g. 100 + """ + + progress_token = ( + self.request_context.meta.progressToken + if self.request_context.meta + else None + ) + + if progress_token is None: + return + + await self.request_context.session.send_progress_notification( + progress_token=progress_token, progress=progress, total=total + ) + + async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]: + """Read a resource by URI. + + Args: + uri: Resource URI to read + + Returns: + The resource content as either text or bytes + """ + assert self._fastmcp is not None, ( + "Context is not available outside of a request" + ) + return await self._fastmcp.read_resource(uri) + + async def log( + self, + level: Literal["debug", "info", "warning", "error"], + message: str, + *, + logger_name: str | None = None, + ) -> None: + """Send a log message to the client. + + Args: + level: Log level (debug, info, warning, error) + message: Log message + logger_name: Optional logger name + **extra: Additional structured data to include + """ + await self.request_context.session.send_log_message( + level=level, data=message, logger=logger_name + ) + + @property + def client_id(self) -> str | None: + """Get the client ID if available.""" + return ( + getattr(self.request_context.meta, "client_id", None) + if self.request_context.meta + else None + ) + + @property + def request_id(self) -> str: + """Get the unique ID for this request.""" + return str(self.request_context.request_id) + + @property + def session(self): + """Access to the underlying session for advanced usage.""" + return self.request_context.session + + # Convenience methods for common log levels + async def debug(self, message: str, **extra: Any) -> None: + """Send a debug log message.""" + await self.log("debug", message, **extra) + + async def info(self, message: str, **extra: Any) -> None: + """Send an info log message.""" + await self.log("info", message, **extra) + + async def warning(self, message: str, **extra: Any) -> None: + """Send a warning log message.""" + await self.log("warning", message, **extra) + + async def error(self, message: str, **extra: Any) -> None: + """Send an error log message.""" + await self.log("error", message, **extra) async def sample( self, diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 4dda5e1d7..2764aa943 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -1,15 +1,15 @@ from typing import Any, cast import mcp.types -from mcp.server.fastmcp.prompts import Prompt -from mcp.server.fastmcp.resources import Resource, ResourceTemplate -from mcp.server.fastmcp.tools.base import Tool -from mcp.server.fastmcp.utilities.func_metadata import func_metadata from mcp.types import BlobResourceContents, PromptMessage, TextResourceContents from fastmcp.clients.base import BaseClient +from fastmcp.prompts import Prompt +from fastmcp.resources import Resource, ResourceTemplate from fastmcp.server.context import Context from fastmcp.server.server import FastMCP +from fastmcp.tools.base import Tool +from fastmcp.utilities.func_metadata import func_metadata from fastmcp.utilities.logging import get_logger logger = get_logger(__name__) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 7d63f6d61..fe1cb0778 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -1,28 +1,93 @@ -from typing import TYPE_CHECKING, Any, Dict +"""FastMCP - A more ergonomic interface for MCP servers.""" -import mcp.server.fastmcp -import mcp.types +from __future__ import annotations as _annotations -from fastmcp.prompts.prompt_manager import PromptManager -from fastmcp.resources.resource_manager import ResourceManager -from fastmcp.server.context import Context -from fastmcp.tools.tool_manager import ToolManager -from fastmcp.utilities.logging import get_logger +import inspect +import json +import re +from collections.abc import AsyncIterator, Callable, Iterable, Sequence +from contextlib import ( + AbstractAsyncContextManager, + asynccontextmanager, +) +from itertools import chain +from typing import TYPE_CHECKING, Any, Generic, Literal + +import anyio +import pydantic_core +import uvicorn +from mcp.server.lowlevel.helper_types import ReadResourceContents +from mcp.server.lowlevel.server import LifespanResultT +from mcp.server.lowlevel.server import Server as MCPServer +from mcp.server.lowlevel.server import lifespan as default_lifespan +from mcp.server.session import ServerSession +from mcp.server.sse import SseServerTransport +from mcp.server.stdio import stdio_server +from mcp.types import ( + AnyFunction, + EmbeddedResource, + GetPromptResult, + ImageContent, + TextContent, +) +from mcp.types import Prompt as MCPPrompt +from mcp.types import PromptArgument as MCPPromptArgument +from mcp.types import Resource as MCPResource +from mcp.types import ResourceTemplate as MCPResourceTemplate +from mcp.types import Tool as MCPTool +from pydantic.networks import AnyUrl +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.routing import Mount, Route + +import fastmcp +import fastmcp.settings +from fastmcp.exceptions import ResourceError +from fastmcp.prompts import Prompt, PromptManager +from fastmcp.resources import FunctionResource, Resource, ResourceManager +from fastmcp.tools import ToolManager +from fastmcp.utilities.logging import configure_logging, get_logger +from fastmcp.utilities.types import Image if TYPE_CHECKING: from fastmcp.clients.base import BaseClient - - from .proxy import FastMCPProxy + from fastmcp.server.context import Context + from fastmcp.server.proxy import FastMCPProxy logger = get_logger(__name__) -class FastMCP(mcp.server.fastmcp.FastMCP): - def __init__(self, name: str | None = None, **settings: Any): - # First initialize with default settings - super().__init__(name=name or "FastMCP", **settings) +def lifespan_wrapper( + app: FastMCP, + lifespan: Callable[[FastMCP], AbstractAsyncContextManager[LifespanResultT]], +) -> Callable[ + [MCPServer[LifespanResultT]], AbstractAsyncContextManager[LifespanResultT] +]: + @asynccontextmanager + async def wrap(s: MCPServer[LifespanResultT]) -> AsyncIterator[LifespanResultT]: + async with lifespan(app) as context: + yield context - # Replace the default managers with our extended ones + return wrap + + +class FastMCP(Generic[LifespanResultT]): + def __init__( + self, + name: str | None = None, + instructions: str | None = None, + lifespan: ( + Callable[["FastMCP"], AbstractAsyncContextManager[LifespanResultT]] | None + ) = None, + **settings: Any, + ): + self.settings = fastmcp.settings.ServerSettings(**settings) + + self._mcp_server = MCPServer[LifespanResultT]( + name=name or "FastMCP", + instructions=instructions, + lifespan=lifespan_wrapper(self, lifespan) if lifespan else default_lifespan, # type: ignore + ) self._tool_manager = ToolManager( warn_on_duplicate_tools=self.settings.warn_on_duplicate_tools ) @@ -32,21 +97,415 @@ class FastMCP(mcp.server.fastmcp.FastMCP): self._prompt_manager = PromptManager( warn_on_duplicate_prompts=self.settings.warn_on_duplicate_prompts ) + self.dependencies = self.settings.dependencies # Setup for mounted apps - self._mounted_apps: Dict[str, "FastMCP"] = {} + self._mounted_apps: dict[str, "FastMCP"] = {} - def get_context(self) -> Context: + # Set up MCP protocol handlers + self._setup_handlers() + + # Configure logging + configure_logging(self.settings.log_level) + + @property + def name(self) -> str: + return self._mcp_server.name + + @property + def instructions(self) -> str | None: + return self._mcp_server.instructions + + def run(self, transport: Literal["stdio", "sse"] = "stdio") -> None: + """Run the FastMCP server. Note this is a synchronous function. + + Args: + transport: Transport protocol to use ("stdio" or "sse") + """ + TRANSPORTS = Literal["stdio", "sse"] + if transport not in TRANSPORTS.__args__: # type: ignore + raise ValueError(f"Unknown transport: {transport}") + + if transport == "stdio": + anyio.run(self.run_stdio_async) + else: # transport == "sse" + anyio.run(self.run_sse_async) + + def _setup_handlers(self) -> None: + """Set up core MCP protocol handlers.""" + self._mcp_server.list_tools()(self.list_tools) + self._mcp_server.call_tool()(self.call_tool) + self._mcp_server.list_resources()(self.list_resources) + self._mcp_server.read_resource()(self.read_resource) + self._mcp_server.list_prompts()(self.list_prompts) + self._mcp_server.get_prompt()(self.get_prompt) + self._mcp_server.list_resource_templates()(self.list_resource_templates) + + async def list_tools(self) -> list[MCPTool]: + """List all available tools.""" + tools = self._tool_manager.list_tools() + return [ + MCPTool( + name=info.name, + description=info.description, + inputSchema=info.parameters, + ) + for info in tools + ] + + def get_context(self) -> "Context[ServerSession, LifespanResultT]": """ Returns a Context object. Note that the context will only be valid during a request; outside a request, most methods will error. """ + try: request_context = self._mcp_server.request_context except LookupError: request_context = None + from fastmcp.server.context import Context + return Context(request_context=request_context, fastmcp=self) + async def call_tool( + self, name: str, arguments: dict[str, Any] + ) -> Sequence[TextContent | ImageContent | EmbeddedResource]: + """Call a tool by name with arguments.""" + context = self.get_context() + result = await self._tool_manager.call_tool(name, arguments, context=context) + converted_result = _convert_to_content(result) + return converted_result + + async def list_resources(self) -> list[MCPResource]: + """List all available resources.""" + + resources = self._resource_manager.list_resources() + return [ + MCPResource( + uri=resource.uri, + name=resource.name or "", + description=resource.description, + mimeType=resource.mime_type, + ) + for resource in resources + ] + + async def list_resource_templates(self) -> list[MCPResourceTemplate]: + templates = self._resource_manager.list_templates() + return [ + MCPResourceTemplate( + uriTemplate=template.uri_template, + name=template.name, + description=template.description, + ) + for template in templates + ] + + async def read_resource(self, uri: AnyUrl | str) -> Iterable[ReadResourceContents]: + """Read a resource by URI.""" + + resource = await self._resource_manager.get_resource(uri) + if not resource: + raise ResourceError(f"Unknown resource: {uri}") + + try: + content = await resource.read() + return [ReadResourceContents(content=content, mime_type=resource.mime_type)] + except Exception as e: + logger.error(f"Error reading resource {uri}: {e}") + raise ResourceError(str(e)) + + def add_tool( + self, + fn: AnyFunction, + name: str | None = None, + description: str | None = None, + ) -> None: + """Add a tool to the server. + + The tool function can optionally request a Context object by adding a parameter + with the Context type annotation. See the @tool decorator for examples. + + Args: + fn: The function to register as a tool + name: Optional name for the tool (defaults to function name) + description: Optional description of what the tool does + """ + self._tool_manager.add_tool(fn, name=name, description=description) + + def tool( + self, name: str | None = None, description: str | None = None + ) -> Callable[[AnyFunction], AnyFunction]: + """Decorator to register a tool. + + Tools can optionally request a Context object by adding a parameter with the + Context type annotation. The context provides access to MCP capabilities like + logging, progress reporting, and resource access. + + Args: + name: Optional name for the tool (defaults to function name) + description: Optional description of what the tool does + + Example: + @server.tool() + def my_tool(x: int) -> str: + return str(x) + + @server.tool() + def tool_with_context(x: int, ctx: Context) -> str: + ctx.info(f"Processing {x}") + return str(x) + + @server.tool() + async def async_tool(x: int, context: Context) -> str: + await context.report_progress(50, 100) + return str(x) + """ + # Check if user passed function directly instead of calling decorator + if callable(name): + raise TypeError( + "The @tool decorator was used incorrectly. " + "Did you forget to call it? Use @tool() instead of @tool" + ) + + def decorator(fn: AnyFunction) -> AnyFunction: + self.add_tool(fn, name=name, description=description) + return fn + + return decorator + + def add_resource(self, resource: Resource) -> None: + """Add a resource to the server. + + Args: + resource: A Resource instance to add + """ + self._resource_manager.add_resource(resource) + + def resource( + self, + uri: str, + *, + name: str | None = None, + description: str | None = None, + mime_type: str | None = None, + ) -> Callable[[AnyFunction], AnyFunction]: + """Decorator to register a function as a resource. + + The function will be called when the resource is read to generate its content. + The function can return: + - str for text content + - bytes for binary content + - other types will be converted to JSON + + If the URI contains parameters (e.g. "resource://{param}") or the function + has parameters, it will be registered as a template resource. + + Args: + uri: URI for the resource (e.g. "resource://my-resource" or "resource://{param}") + name: Optional name for the resource + description: Optional description of the resource + mime_type: Optional MIME type for the resource + + Example: + @server.resource("resource://my-resource") + def get_data() -> str: + return "Hello, world!" + + @server.resource("resource://my-resource") + async get_data() -> str: + data = await fetch_data() + return f"Hello, world! {data}" + + @server.resource("resource://{city}/weather") + def get_weather(city: str) -> str: + return f"Weather for {city}" + + @server.resource("resource://{city}/weather") + async def get_weather(city: str) -> str: + data = await fetch_weather(city) + return f"Weather for {city}: {data}" + """ + # Check if user passed function directly instead of calling decorator + if callable(uri): + raise TypeError( + "The @resource decorator was used incorrectly. " + "Did you forget to call it? Use @resource('uri') instead of @resource" + ) + + def decorator(fn: AnyFunction) -> AnyFunction: + # Check if this should be a template + has_uri_params = "{" in uri and "}" in uri + has_func_params = bool(inspect.signature(fn).parameters) + + if has_uri_params or has_func_params: + # Validate that URI params match function params + uri_params = set(re.findall(r"{(\w+)}", uri)) + func_params = set(inspect.signature(fn).parameters.keys()) + + if uri_params != func_params: + raise ValueError( + f"Mismatch between URI parameters {uri_params} " + f"and function parameters {func_params}" + ) + + # Register as template + self._resource_manager.add_template( + fn=fn, + uri_template=uri, + name=name, + description=description, + mime_type=mime_type or "text/plain", + ) + else: + # Register as regular resource + resource = FunctionResource( + uri=AnyUrl(uri), + name=name, + description=description, + mime_type=mime_type or "text/plain", + fn=fn, + ) + self.add_resource(resource) + return fn + + return decorator + + def add_prompt(self, prompt: Prompt) -> None: + """Add a prompt to the server. + + Args: + prompt: A Prompt instance to add + """ + self._prompt_manager.add_prompt(prompt) + + def prompt( + self, name: str | None = None, description: str | None = None + ) -> Callable[[AnyFunction], AnyFunction]: + """Decorator to register a prompt. + + Args: + name: Optional name for the prompt (defaults to function name) + description: Optional description of what the prompt does + + Example: + @server.prompt() + def analyze_table(table_name: str) -> list[Message]: + schema = read_table_schema(table_name) + return [ + { + "role": "user", + "content": f"Analyze this schema:\n{schema}" + } + ] + + @server.prompt() + async def analyze_file(path: str) -> list[Message]: + content = await read_file(path) + return [ + { + "role": "user", + "content": { + "type": "resource", + "resource": { + "uri": f"file://{path}", + "text": content + } + } + } + ] + """ + # Check if user passed function directly instead of calling decorator + if callable(name): + raise TypeError( + "The @prompt decorator was used incorrectly. " + "Did you forget to call it? Use @prompt() instead of @prompt" + ) + + def decorator(func: AnyFunction) -> AnyFunction: + prompt = Prompt.from_function(func, name=name, description=description) + self.add_prompt(prompt) + return func + + return decorator + + async def run_stdio_async(self) -> None: + """Run the server using stdio transport.""" + async with stdio_server() as (read_stream, write_stream): + await self._mcp_server.run( + read_stream, + write_stream, + self._mcp_server.create_initialization_options(), + ) + + async def run_sse_async(self) -> None: + """Run the server using SSE transport.""" + starlette_app = self.sse_app() + + config = uvicorn.Config( + starlette_app, + host=self.settings.host, + port=self.settings.port, + log_level=self.settings.log_level.lower(), + ) + server = uvicorn.Server(config) + await server.serve() + + def sse_app(self) -> Starlette: + """Return an instance of the SSE server app.""" + sse = SseServerTransport(self.settings.message_path) + + async def handle_sse(request: Request) -> None: + async with sse.connect_sse( + request.scope, + request.receive, + request._send, # type: ignore[reportPrivateUsage] + ) as streams: + await self._mcp_server.run( + streams[0], + streams[1], + self._mcp_server.create_initialization_options(), + ) + + return Starlette( + debug=self.settings.debug, + routes=[ + Route(self.settings.sse_path, endpoint=handle_sse), + Mount(self.settings.message_path, app=sse.handle_post_message), + ], + ) + + async def list_prompts(self) -> list[MCPPrompt]: + """List all available prompts.""" + prompts = self._prompt_manager.list_prompts() + return [ + MCPPrompt( + name=prompt.name, + description=prompt.description, + arguments=[ + MCPPromptArgument( + name=arg.name, + description=arg.description, + required=arg.required, + ) + for arg in (prompt.arguments or []) + ], + ) + for prompt in prompts + ] + + async def get_prompt( + self, name: str, arguments: dict[str, Any] | None = None + ) -> GetPromptResult: + """Get a prompt by name with arguments.""" + try: + messages = await self._prompt_manager.render_prompt(name, arguments) + + return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages)) + except Exception as e: + logger.error(f"Error getting prompt {name}: {e}") + raise ValueError(str(e)) + def mount(self, prefix: str, app: "FastMCP") -> None: """Mount another FastMCP application with a given prefix. @@ -105,3 +564,28 @@ class FastMCP(mcp.server.fastmcp.FastMCP): from .proxy import FastMCPProxy return await FastMCPProxy.from_client(client=client, **settings) + + +def _convert_to_content( + result: Any, +) -> Sequence[TextContent | ImageContent | EmbeddedResource]: + """Convert a result to a sequence of content objects.""" + if result is None: + return [] + + if isinstance(result, TextContent | ImageContent | EmbeddedResource): + return [result] + + if isinstance(result, Image): + return [result.to_image_content()] + + if isinstance(result, list | tuple): + return list(chain.from_iterable(_convert_to_content(item) for item in result)) # type: ignore[reportUnknownVariableType] + + if not isinstance(result, str): + try: + result = json.dumps(pydantic_core.to_jsonable_python(result)) + except Exception: + result = str(result) + + return [TextContent(type="text", text=result)] diff --git a/src/fastmcp/settings.py b/src/fastmcp/settings.py index c89bde116..1561a39a2 100644 --- a/src/fastmcp/settings.py +++ b/src/fastmcp/settings.py @@ -1,25 +1,73 @@ -from typing import Literal +from __future__ import annotations as _annotations +from typing import TYPE_CHECKING, Literal + +from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict +if TYPE_CHECKING: + pass + LOG_LEVEL = Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] class Settings(BaseSettings): + """FastMCP settings.""" + + model_config = SettingsConfigDict( + env_prefix="FASTMCP_", + env_file=".env", + extra="ignore", + ) + + test_mode: bool = False + log_level: LOG_LEVEL = "INFO" + + +class ServerSettings(BaseSettings): """FastMCP server settings. All settings can be configured via environment variables with the prefix FASTMCP_. For example, FASTMCP_DEBUG=true will set debug=True. """ - model_config: SettingsConfigDict = SettingsConfigDict( - env_prefix="FASTMCP_", + model_config = SettingsConfigDict( + env_prefix="FASTMCP_SERVER_", env_file=".env", extra="ignore", ) - debug: bool = False - log_level: LOG_LEVEL = "INFO" + log_level: LOG_LEVEL = Field(default_factory=lambda: Settings().log_level) - # Client settings - client_log_level: LOG_LEVEL | None = None + # HTTP settings + host: str = "0.0.0.0" + port: int = 8000 + sse_path: str = "/sse" + message_path: str = "/messages/" + debug: bool = False + + # resource settings + warn_on_duplicate_resources: bool = True + + # tool settings + warn_on_duplicate_tools: bool = True + + # prompt settings + warn_on_duplicate_prompts: bool = True + + dependencies: list[str] = Field( + default_factory=list, + description="List of dependencies to install in the server environment", + ) + + +class ClientSettings(BaseSettings): + """FastMCP client settings.""" + + model_config = SettingsConfigDict( + env_prefix="FASTMCP_CLIENT_", + env_file=".env", + extra="ignore", + ) + + log_level: LOG_LEVEL = Field(default_factory=lambda: Settings().log_level) diff --git a/src/fastmcp/tools/__init__.py b/src/fastmcp/tools/__init__.py index e69de29bb..ae9c65619 100644 --- a/src/fastmcp/tools/__init__.py +++ b/src/fastmcp/tools/__init__.py @@ -0,0 +1,4 @@ +from .base import Tool +from .tool_manager import ToolManager + +__all__ = ["Tool", "ToolManager"] diff --git a/src/fastmcp/tools/base.py b/src/fastmcp/tools/base.py new file mode 100644 index 000000000..27305cb5a --- /dev/null +++ b/src/fastmcp/tools/base.py @@ -0,0 +1,93 @@ +from __future__ import annotations as _annotations + +import inspect +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from pydantic import BaseModel, Field + +from fastmcp.exceptions import ToolError +from fastmcp.utilities.func_metadata import FuncMetadata, func_metadata + +if TYPE_CHECKING: + from mcp.server.session import ServerSessionT + from mcp.shared.context import LifespanContextT + + from fastmcp.server import Context + + +class Tool(BaseModel): + """Internal tool registration info.""" + + fn: Callable[..., Any] = Field(exclude=True) + name: str = Field(description="Name of the tool") + description: str = Field(description="Description of what the tool does") + parameters: dict[str, Any] = Field(description="JSON schema for tool parameters") + fn_metadata: FuncMetadata = Field( + description="Metadata about the function including a pydantic model for tool" + " arguments" + ) + is_async: bool = Field(description="Whether the tool is async") + context_kwarg: str | None = Field( + None, description="Name of the kwarg that should receive context" + ) + + @classmethod + def from_function( + cls, + fn: Callable[..., Any], + name: str | None = None, + description: str | None = None, + context_kwarg: str | None = None, + ) -> Tool: + """Create a Tool from a function.""" + from fastmcp import Context + + func_name = name or fn.__name__ + + if func_name == "": + raise ValueError("You must provide a name for lambda functions") + + func_doc = description or fn.__doc__ or "" + is_async = inspect.iscoroutinefunction(fn) + + if context_kwarg is None: + sig = inspect.signature(fn) + for param_name, param in sig.parameters.items(): + if param.annotation is Context: + context_kwarg = param_name + break + + func_arg_metadata = func_metadata( + fn, + skip_names=[context_kwarg] if context_kwarg is not None else [], + ) + parameters = func_arg_metadata.arg_model.model_json_schema() + + return cls( + fn=fn, + name=func_name, + description=func_doc, + parameters=parameters, + fn_metadata=func_arg_metadata, + is_async=is_async, + context_kwarg=context_kwarg, + ) + + async def run( + self, + arguments: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, + ) -> Any: + """Run the tool with arguments.""" + try: + return await self.fn_metadata.call_fn_with_arg_validation( + self.fn, + self.is_async, + arguments, + {self.context_kwarg: context} + if self.context_kwarg is not None + else None, + ) + except Exception as e: + raise ToolError(f"Error executing tool {self.name}: {e}") from e diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index 1b3b95533..48e2fdc9a 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -1,16 +1,65 @@ -import mcp.server.fastmcp.tools -from mcp.server.fastmcp.tools import Tool +from __future__ import annotations as _annotations +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from mcp.shared.context import LifespanContextT + +from fastmcp.exceptions import ToolError +from fastmcp.tools.base import Tool from fastmcp.utilities.logging import get_logger +if TYPE_CHECKING: + from mcp.server.session import ServerSessionT + + from fastmcp.server import Context + logger = get_logger(__name__) -class ToolManager(mcp.server.fastmcp.tools.ToolManager): - """ - Extended ToolManager that supports importing tools from other managers. - Adds ability to import tools from other managers with prefixed names. - """ +class ToolManager: + """Manages FastMCP tools.""" + + def __init__(self, warn_on_duplicate_tools: bool = True): + self._tools: dict[str, Tool] = {} + self.warn_on_duplicate_tools = warn_on_duplicate_tools + + def get_tool(self, name: str) -> Tool | None: + """Get tool by name.""" + return self._tools.get(name) + + def list_tools(self) -> list[Tool]: + """List all registered tools.""" + return list(self._tools.values()) + + def add_tool( + self, + fn: Callable[..., Any], + name: str | None = None, + description: str | None = None, + ) -> Tool: + """Add a tool to the server.""" + tool = Tool.from_function(fn, name=name, description=description) + existing = self._tools.get(tool.name) + if existing: + if self.warn_on_duplicate_tools: + logger.warning(f"Tool already exists: {tool.name}") + return existing + self._tools[tool.name] = tool + return tool + + async def call_tool( + self, + name: str, + arguments: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, + ) -> Any: + """Call a tool by name with arguments.""" + tool = self.get_tool(name) + if not tool: + raise ToolError(f"Unknown tool: {name}") + + return await tool.run(arguments, context=context) def import_tools( self, tool_manager: "ToolManager", prefix: str | None = None diff --git a/src/fastmcp/utilities/func_metadata.py b/src/fastmcp/utilities/func_metadata.py new file mode 100644 index 000000000..5673b5a26 --- /dev/null +++ b/src/fastmcp/utilities/func_metadata.py @@ -0,0 +1,214 @@ +import inspect +import json +from collections.abc import Awaitable, Callable, Sequence +from typing import ( + Annotated, + Any, + ForwardRef, +) + +from pydantic import BaseModel, ConfigDict, Field, WithJsonSchema, create_model +from pydantic._internal._typing_extra import eval_type_backport +from pydantic.fields import FieldInfo +from pydantic_core import PydanticUndefined + +from fastmcp.exceptions import InvalidSignature +from fastmcp.utilities.logging import get_logger + +logger = get_logger(__name__) + + +class ArgModelBase(BaseModel): + """A model representing the arguments to a function.""" + + def model_dump_one_level(self) -> dict[str, Any]: + """Return a dict of the model's fields, one level deep. + + That is, sub-models etc are not dumped - they are kept as pydantic models. + """ + kwargs: dict[str, Any] = {} + for field_name in self.model_fields.keys(): + kwargs[field_name] = getattr(self, field_name) + return kwargs + + model_config = ConfigDict( + arbitrary_types_allowed=True, + ) + + +class FuncMetadata(BaseModel): + arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)] + # We can add things in the future like + # - Maybe some args are excluded from attempting to parse from JSON + # - Maybe some args are special (like context) for dependency injection + + async def call_fn_with_arg_validation( + self, + fn: Callable[..., Any] | Awaitable[Any], + fn_is_async: bool, + arguments_to_validate: dict[str, Any], + arguments_to_pass_directly: dict[str, Any] | None, + ) -> Any: + """Call the given function with arguments validated and injected. + + Arguments are first attempted to be parsed from JSON, then validated against + the argument model, before being passed to the function. + """ + arguments_pre_parsed = self.pre_parse_json(arguments_to_validate) + arguments_parsed_model = self.arg_model.model_validate(arguments_pre_parsed) + arguments_parsed_dict = arguments_parsed_model.model_dump_one_level() + + arguments_parsed_dict |= arguments_to_pass_directly or {} + + if fn_is_async: + if isinstance(fn, Awaitable): + return await fn + return await fn(**arguments_parsed_dict) + if isinstance(fn, Callable): + return fn(**arguments_parsed_dict) + raise TypeError("fn must be either Callable or Awaitable") + + def pre_parse_json(self, data: dict[str, Any]) -> dict[str, Any]: + """Pre-parse data from JSON. + + Return a dict with same keys as input but with values parsed from JSON + if appropriate. + + This is to handle cases like `["a", "b", "c"]` being passed in as JSON inside + a string rather than an actual list. Claude desktop is prone to this - in fact + it seems incapable of NOT doing this. For sub-models, it tends to pass + dicts (JSON objects) as JSON strings, which can be pre-parsed here. + """ + new_data = data.copy() # Shallow copy + for field_name, _field_info in self.arg_model.model_fields.items(): + if field_name not in data.keys(): + continue + if isinstance(data[field_name], str): + try: + pre_parsed = json.loads(data[field_name]) + except json.JSONDecodeError: + continue # Not JSON - skip + if isinstance(pre_parsed, str | int | float): + # This is likely that the raw value is e.g. `"hello"` which we + # Should really be parsed as '"hello"' in Python - but if we parse + # it as JSON it'll turn into just 'hello'. So we skip it. + continue + new_data[field_name] = pre_parsed + assert new_data.keys() == data.keys() + return new_data + + model_config = ConfigDict( + arbitrary_types_allowed=True, + ) + + +def func_metadata( + func: Callable[..., Any], skip_names: Sequence[str] = () +) -> FuncMetadata: + """Given a function, return metadata including a pydantic model representing its + signature. + + The use case for this is + ``` + meta = func_to_pyd(func) + validated_args = meta.arg_model.model_validate(some_raw_data_dict) + return func(**validated_args.model_dump_one_level()) + ``` + + **critically** it also provides pre-parse helper to attempt to parse things from + JSON. + + Args: + func: The function to convert to a pydantic model + skip_names: A list of parameter names to skip. These will not be included in + the model. + Returns: + A pydantic model representing the function's signature. + """ + sig = _get_typed_signature(func) + params = sig.parameters + dynamic_pydantic_model_params: dict[str, Any] = {} + globalns = getattr(func, "__globals__", {}) + for param in params.values(): + if param.name.startswith("_"): + raise InvalidSignature( + f"Parameter {param.name} of {func.__name__} cannot start with '_'" + ) + if param.name in skip_names: + continue + annotation = param.annotation + + # `x: None` / `x: None = None` + if annotation is None: + annotation = Annotated[ + None, + Field( + default=param.default + if param.default is not inspect.Parameter.empty + else PydanticUndefined + ), + ] + + # Untyped field + if annotation is inspect.Parameter.empty: + annotation = Annotated[ + Any, + Field(), + # 🤷 + WithJsonSchema({"title": param.name, "type": "string"}), + ] + + field_info = FieldInfo.from_annotated_attribute( + _get_typed_annotation(annotation, globalns), + param.default + if param.default is not inspect.Parameter.empty + else PydanticUndefined, + ) + dynamic_pydantic_model_params[param.name] = (field_info.annotation, field_info) + continue + + arguments_model = create_model( + f"{func.__name__}Arguments", + **dynamic_pydantic_model_params, + __base__=ArgModelBase, + ) + resp = FuncMetadata(arg_model=arguments_model) + return resp + + +def _get_typed_annotation(annotation: Any, globalns: dict[str, Any]) -> Any: + def try_eval_type( + value: Any, globalns: dict[str, Any], localns: dict[str, Any] + ) -> tuple[Any, bool]: + try: + return eval_type_backport(value, globalns, localns), True + except NameError: + return value, False + + if isinstance(annotation, str): + annotation = ForwardRef(annotation) + annotation, status = try_eval_type(annotation, globalns, globalns) + + # This check and raise could perhaps be skipped, and we (FastMCP) just call + # model_rebuild right before using it 🤷 + if status is False: + raise InvalidSignature(f"Unable to evaluate type annotation {annotation}") + + return annotation + + +def _get_typed_signature(call: Callable[..., Any]) -> inspect.Signature: + """Get function signature while evaluating forward references""" + signature = inspect.signature(call) + globalns = getattr(call, "__globals__", {}) + typed_params = [ + inspect.Parameter( + name=param.name, + kind=param.kind, + default=param.default, + annotation=_get_typed_annotation(param.annotation, globalns), + ) + for param in signature.parameters.values() + ] + typed_signature = inspect.Signature(typed_params) + return typed_signature diff --git a/src/fastmcp/utilities/types.py b/src/fastmcp/utilities/types.py index b93d244f0..ccaa3d69a 100644 --- a/src/fastmcp/utilities/types.py +++ b/src/fastmcp/utilities/types.py @@ -2,7 +2,6 @@ import base64 from pathlib import Path -from typing import Optional, Union from mcp.types import ImageContent @@ -12,9 +11,9 @@ class Image: def __init__( self, - path: Optional[Union[str, Path]] = None, - data: Optional[bytes] = None, - format: Optional[str] = None, + path: str | Path | None = None, + data: bytes | None = None, + format: str | None = None, ): if path is None and data is None: raise ValueError("Either path or data must be provided") diff --git a/tests/prompts/__init__.py b/tests/prompts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/prompts/test_base.py b/tests/prompts/test_base.py new file mode 100644 index 000000000..a3b591858 --- /dev/null +++ b/tests/prompts/test_base.py @@ -0,0 +1,206 @@ +import pytest +from mcp.types import EmbeddedResource, TextResourceContents +from pydantic import FileUrl + +from fastmcp.prompts.base import ( + AssistantMessage, + Message, + Prompt, + TextContent, + UserMessage, +) + + +class TestRenderPrompt: + @pytest.mark.anyio + async def test_basic_fn(self): + def fn() -> str: + return "Hello, world!" + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [ + UserMessage(content=TextContent(type="text", text="Hello, world!")) + ] + + @pytest.mark.anyio + async def test_async_fn(self): + async def fn() -> str: + return "Hello, world!" + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [ + UserMessage(content=TextContent(type="text", text="Hello, world!")) + ] + + @pytest.mark.anyio + async def test_fn_with_args(self): + async def fn(name: str, age: int = 30) -> str: + return f"Hello, {name}! You're {age} years old." + + prompt = Prompt.from_function(fn) + assert await prompt.render(arguments=dict(name="World")) == [ + UserMessage( + content=TextContent( + type="text", text="Hello, World! You're 30 years old." + ) + ) + ] + + @pytest.mark.anyio + async def test_fn_with_invalid_kwargs(self): + async def fn(name: str, age: int = 30) -> str: + return f"Hello, {name}! You're {age} years old." + + prompt = Prompt.from_function(fn) + with pytest.raises(ValueError): + await prompt.render(arguments=dict(age=40)) + + @pytest.mark.anyio + async def test_fn_returns_message(self): + async def fn() -> UserMessage: + return UserMessage(content="Hello, world!") + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [ + UserMessage(content=TextContent(type="text", text="Hello, world!")) + ] + + @pytest.mark.anyio + async def test_fn_returns_assistant_message(self): + async def fn() -> AssistantMessage: + return AssistantMessage( + content=TextContent(type="text", text="Hello, world!") + ) + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [ + AssistantMessage(content=TextContent(type="text", text="Hello, world!")) + ] + + @pytest.mark.anyio + async def test_fn_returns_multiple_messages(self): + expected = [ + UserMessage("Hello, world!"), + AssistantMessage("How can I help you today?"), + UserMessage("I'm looking for a restaurant in the center of town."), + ] + + async def fn() -> list[Message]: + return expected + + prompt = Prompt.from_function(fn) + assert await prompt.render() == expected + + @pytest.mark.anyio + async def test_fn_returns_list_of_strings(self): + expected = [ + "Hello, world!", + "I'm looking for a restaurant in the center of town.", + ] + + async def fn() -> list[str]: + return expected + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [UserMessage(t) for t in expected] + + @pytest.mark.anyio + async def test_fn_returns_resource_content(self): + """Test returning a message with resource content.""" + + async def fn() -> UserMessage: + return UserMessage( + content=EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri=FileUrl("file://file.txt"), + text="File contents", + mimeType="text/plain", + ), + ) + ) + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [ + UserMessage( + content=EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri=FileUrl("file://file.txt"), + text="File contents", + mimeType="text/plain", + ), + ) + ) + ] + + @pytest.mark.anyio + async def test_fn_returns_mixed_content(self): + """Test returning messages with mixed content types.""" + + async def fn() -> list[Message]: + return [ + UserMessage(content="Please analyze this file:"), + UserMessage( + content=EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri=FileUrl("file://file.txt"), + text="File contents", + mimeType="text/plain", + ), + ) + ), + AssistantMessage(content="I'll help analyze that file."), + ] + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [ + UserMessage( + content=TextContent(type="text", text="Please analyze this file:") + ), + UserMessage( + content=EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri=FileUrl("file://file.txt"), + text="File contents", + mimeType="text/plain", + ), + ) + ), + AssistantMessage( + content=TextContent(type="text", text="I'll help analyze that file.") + ), + ] + + @pytest.mark.anyio + async def test_fn_returns_dict_with_resource(self): + """Test returning a dict with resource content.""" + + async def fn() -> dict: + return { + "role": "user", + "content": { + "type": "resource", + "resource": { + "uri": FileUrl("file://file.txt"), + "text": "File contents", + "mimeType": "text/plain", + }, + }, + } + + prompt = Prompt.from_function(fn) + assert await prompt.render() == [ + UserMessage( + content=EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri=FileUrl("file://file.txt"), + text="File contents", + mimeType="text/plain", + ), + ) + ) + ] diff --git a/tests/prompts/test_prompt_manager.py b/tests/prompts/test_prompt_manager.py index 3a5accec0..1964db74d 100644 --- a/tests/prompts/test_prompt_manager.py +++ b/tests/prompts/test_prompt_manager.py @@ -1,166 +1,283 @@ -from mcp.server.fastmcp.prompts import Prompt -from mcp.server.fastmcp.prompts.base import PromptArgument +import pytest +from fastmcp.prompts import Prompt +from fastmcp.prompts.base import PromptArgument, TextContent, UserMessage from fastmcp.prompts.prompt_manager import PromptManager -def test_import_prompts(): - """Test importing prompts from one manager to another with a prefix.""" - # Setup source manager with prompts - source_manager = PromptManager() +class TestPromptManager: + def test_add_prompt(self): + """Test adding a prompt to the manager.""" - # Create test prompts with proper function handlers - async def summary_fn(**kwargs): - return [{"role": "assistant", "content": f"Summary of: {kwargs.get('text')}"}] + def fn() -> str: + return "Hello, world!" - async def translate_fn(**kwargs): - return [ - { - "role": "assistant", - "content": f"Translation to {kwargs.get('language')}: {kwargs.get('text')}", - } + manager = PromptManager() + prompt = Prompt.from_function(fn) + added = manager.add_prompt(prompt) + assert added == prompt + assert manager.get_prompt("fn") == prompt + + def test_add_duplicate_prompt(self, caplog): + """Test adding the same prompt twice.""" + + def fn() -> str: + return "Hello, world!" + + manager = PromptManager() + prompt = Prompt.from_function(fn) + first = manager.add_prompt(prompt) + second = manager.add_prompt(prompt) + assert first == second + assert "Prompt already exists" in caplog.text + + def test_disable_warn_on_duplicate_prompts(self, caplog): + """Test disabling warning on duplicate prompts.""" + + def fn() -> str: + return "Hello, world!" + + manager = PromptManager(warn_on_duplicate_prompts=False) + prompt = Prompt.from_function(fn) + first = manager.add_prompt(prompt) + second = manager.add_prompt(prompt) + assert first == second + assert "Prompt already exists" not in caplog.text + + def test_list_prompts(self): + """Test listing all prompts.""" + + def fn1() -> str: + return "Hello, world!" + + def fn2() -> str: + return "Goodbye, world!" + + manager = PromptManager() + prompt1 = Prompt.from_function(fn1) + prompt2 = Prompt.from_function(fn2) + manager.add_prompt(prompt1) + manager.add_prompt(prompt2) + prompts = manager.list_prompts() + assert len(prompts) == 2 + assert prompts == [prompt1, prompt2] + + @pytest.mark.anyio + async def test_render_prompt(self): + """Test rendering a prompt.""" + + def fn() -> str: + return "Hello, world!" + + manager = PromptManager() + prompt = Prompt.from_function(fn) + manager.add_prompt(prompt) + messages = await manager.render_prompt("fn") + assert messages == [ + UserMessage(content=TextContent(type="text", text="Hello, world!")) ] - summary_prompt = Prompt( - name="summary", - description="Generate a summary of text", - arguments=[PromptArgument(name="text", description="Text to summarize")], - fn=summary_fn, - ) - source_manager._prompts["summary"] = summary_prompt + @pytest.mark.anyio + async def test_render_prompt_with_args(self): + """Test rendering a prompt with arguments.""" - translate_prompt = Prompt( - name="translate", - description="Translate text to another language", - arguments=[ - PromptArgument(name="text", description="Text to translate"), - PromptArgument(name="language", description="Target language"), - ], - fn=translate_fn, - ) - source_manager._prompts["translate"] = translate_prompt + def fn(name: str) -> str: + return f"Hello, {name}!" - # Create target manager - target_manager = PromptManager() + manager = PromptManager() + prompt = Prompt.from_function(fn) + manager.add_prompt(prompt) + messages = await manager.render_prompt("fn", arguments={"name": "World"}) + assert messages == [ + UserMessage(content=TextContent(type="text", text="Hello, World!")) + ] - # Import prompts from source to target - prefix = "nlp/" - target_manager.import_prompts(source_manager, prefix) + @pytest.mark.anyio + async def test_render_unknown_prompt(self): + """Test rendering a non-existent prompt.""" + manager = PromptManager() + with pytest.raises(ValueError, match="Unknown prompt: unknown"): + await manager.render_prompt("unknown") - # Verify prompts were imported with prefixes - assert "nlp/summary" in target_manager._prompts - assert "nlp/translate" in target_manager._prompts + @pytest.mark.anyio + async def test_render_prompt_with_missing_args(self): + """Test rendering a prompt with missing required arguments.""" - # Verify the original prompts still exist in source manager - assert "summary" in source_manager._prompts - assert "translate" in source_manager._prompts + def fn(name: str) -> str: + return f"Hello, {name}!" - # Verify the imported prompts have the correct properties - assert target_manager._prompts["nlp/summary"].name == "summary" - assert ( - target_manager._prompts["nlp/summary"].description - == "Generate a summary of text" - ) + manager = PromptManager() + prompt = Prompt.from_function(fn) + manager.add_prompt(prompt) + with pytest.raises(ValueError, match="Missing required arguments"): + await manager.render_prompt("fn") - assert target_manager._prompts["nlp/translate"].name == "translate" - assert ( - target_manager._prompts["nlp/translate"].description - == "Translate text to another language" - ) - # Verify functions were properly copied - if hasattr(target_manager._prompts["nlp/summary"], "fn"): - assert target_manager._prompts["nlp/summary"].fn.__name__ == summary_fn.__name__ +class TestImports: + def test_import_prompts(self): + """Test importing prompts from one manager to another with a prefix.""" + # Setup source manager with prompts + source_manager = PromptManager() - if hasattr(target_manager._prompts["nlp/translate"], "fn"): + # Create test prompts with proper function handlers + async def summary_fn(**kwargs): + return [ + {"role": "assistant", "content": f"Summary of: {kwargs.get('text')}"} + ] + + async def translate_fn(**kwargs): + return [ + { + "role": "assistant", + "content": f"Translation to {kwargs.get('language')}: {kwargs.get('text')}", + } + ] + + summary_prompt = Prompt( + name="summary", + description="Generate a summary of text", + arguments=[PromptArgument(name="text", description="Text to summarize")], + fn=summary_fn, + ) + source_manager._prompts["summary"] = summary_prompt + + translate_prompt = Prompt( + name="translate", + description="Translate text to another language", + arguments=[ + PromptArgument(name="text", description="Text to translate"), + PromptArgument(name="language", description="Target language"), + ], + fn=translate_fn, + ) + source_manager._prompts["translate"] = translate_prompt + + # Create target manager + target_manager = PromptManager() + + # Import prompts from source to target + prefix = "nlp/" + target_manager.import_prompts(source_manager, prefix) + + # Verify prompts were imported with prefixes + assert "nlp/summary" in target_manager._prompts + assert "nlp/translate" in target_manager._prompts + + # Verify the original prompts still exist in source manager + assert "summary" in source_manager._prompts + assert "translate" in source_manager._prompts + + # Verify the imported prompts have the correct properties + assert target_manager._prompts["nlp/summary"].name == "summary" assert ( - target_manager._prompts["nlp/translate"].fn.__name__ - == translate_fn.__name__ + target_manager._prompts["nlp/summary"].description + == "Generate a summary of text" ) - -def test_import_prompts_with_duplicates(): - """Test handling of duplicate prompts during import.""" - # Setup source and target managers with same prompt names - source_manager = PromptManager() - target_manager = PromptManager() - - # Add the same prompt name to both managers with functions - async def source_fn(**kwargs): - return [{"role": "assistant", "content": "Source content"}] - - async def target_fn(**kwargs): - return [{"role": "assistant", "content": "Target content"}] - - source_prompt = Prompt( - name="common", - description="Source description", - arguments=None, - fn=source_fn, - ) - source_manager._prompts["common"] = source_prompt - - target_prompt = Prompt( - name="common", - description="Target description", - arguments=None, - fn=target_fn, - ) - target_manager._prompts["common"] = target_prompt - - # Import prompts with prefix - prefix = "external/" - target_manager.import_prompts(source_manager, prefix) - - # Verify both prompts exist in target manager - assert "common" in target_manager._prompts - assert "external/common" in target_manager._prompts - - # Verify the functions of both prompts - if hasattr(target_manager._prompts["common"], "fn") and hasattr( - target_manager._prompts["external/common"], "fn" - ): - assert target_manager._prompts["common"].fn.__name__ == target_fn.__name__ + assert target_manager._prompts["nlp/translate"].name == "translate" assert ( - target_manager._prompts["external/common"].fn.__name__ == source_fn.__name__ + target_manager._prompts["nlp/translate"].description + == "Translate text to another language" ) + # Verify functions were properly copied + if hasattr(target_manager._prompts["nlp/summary"], "fn"): + assert ( + target_manager._prompts["nlp/summary"].fn.__name__ + == summary_fn.__name__ + ) -def test_import_prompts_with_nested_prefixes(): - """Test importing already prefixed prompts.""" - # Setup source manager with already prefixed prompts - first_manager = PromptManager() - second_manager = PromptManager() - third_manager = PromptManager() + if hasattr(target_manager._prompts["nlp/translate"], "fn"): + assert ( + target_manager._prompts["nlp/translate"].fn.__name__ + == translate_fn.__name__ + ) - # Add prompt to first manager with a function - async def analyze_fn(**kwargs): - return [{"role": "assistant", "content": f"Analysis of: {kwargs.get('text')}"}] + def test_import_prompts_with_duplicates(self): + """Test handling of duplicate prompts during import.""" + # Setup source and target managers with same prompt names + source_manager = PromptManager() + target_manager = PromptManager() - original_prompt = Prompt( - name="analyze", - description="Analyze text", - arguments=[PromptArgument(name="text", description="Text to analyze")], - fn=analyze_fn, - ) - first_manager._prompts["analyze"] = original_prompt + # Add the same prompt name to both managers with functions + async def source_fn(**kwargs): + return [{"role": "assistant", "content": "Source content"}] - # Import to second manager with prefix - second_manager.import_prompts(first_manager, "text/") + async def target_fn(**kwargs): + return [{"role": "assistant", "content": "Target content"}] - # Import from second to third with another prefix - third_manager.import_prompts(second_manager, "ai/") - - # Verify the nested prefixing - assert "text/analyze" in second_manager._prompts - assert "ai/text/analyze" in third_manager._prompts - - # Verify the properties of the most nested prompt - assert third_manager._prompts["ai/text/analyze"].name == "analyze" - assert third_manager._prompts["ai/text/analyze"].description == "Analyze text" - - # Verify function was properly copied through multiple imports - if hasattr(third_manager._prompts["ai/text/analyze"], "fn"): - assert ( - third_manager._prompts["ai/text/analyze"].fn.__name__ == analyze_fn.__name__ + source_prompt = Prompt( + name="common", + description="Source description", + arguments=None, + fn=source_fn, ) + source_manager._prompts["common"] = source_prompt + + target_prompt = Prompt( + name="common", + description="Target description", + arguments=None, + fn=target_fn, + ) + target_manager._prompts["common"] = target_prompt + + # Import prompts with prefix + prefix = "external/" + target_manager.import_prompts(source_manager, prefix) + + # Verify both prompts exist in target manager + assert "common" in target_manager._prompts + assert "external/common" in target_manager._prompts + + # Verify the functions of both prompts + if hasattr(target_manager._prompts["common"], "fn") and hasattr( + target_manager._prompts["external/common"], "fn" + ): + assert target_manager._prompts["common"].fn.__name__ == target_fn.__name__ + assert ( + target_manager._prompts["external/common"].fn.__name__ + == source_fn.__name__ + ) + + def test_import_prompts_with_nested_prefixes(self): + """Test importing already prefixed prompts.""" + # Setup source manager with already prefixed prompts + first_manager = PromptManager() + second_manager = PromptManager() + third_manager = PromptManager() + + # Add prompt to first manager with a function + async def analyze_fn(**kwargs): + return [ + {"role": "assistant", "content": f"Analysis of: {kwargs.get('text')}"} + ] + + original_prompt = Prompt( + name="analyze", + description="Analyze text", + arguments=[PromptArgument(name="text", description="Text to analyze")], + fn=analyze_fn, + ) + first_manager._prompts["analyze"] = original_prompt + + # Import to second manager with prefix + second_manager.import_prompts(first_manager, "text/") + + # Import from second to third with another prefix + third_manager.import_prompts(second_manager, "ai/") + + # Verify the nested prefixing + assert "text/analyze" in second_manager._prompts + assert "ai/text/analyze" in third_manager._prompts + + # Verify the properties of the most nested prompt + assert third_manager._prompts["ai/text/analyze"].name == "analyze" + assert third_manager._prompts["ai/text/analyze"].description == "Analyze text" + + # Verify function was properly copied through multiple imports + if hasattr(third_manager._prompts["ai/text/analyze"], "fn"): + assert ( + third_manager._prompts["ai/text/analyze"].fn.__name__ + == analyze_fn.__name__ + ) diff --git a/tests/resources/__init__.py b/tests/resources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/resources/test_file_resources.py b/tests/resources/test_file_resources.py new file mode 100644 index 000000000..ef05e32f3 --- /dev/null +++ b/tests/resources/test_file_resources.py @@ -0,0 +1,119 @@ +import os +from pathlib import Path +from tempfile import NamedTemporaryFile + +import pytest +from pydantic import FileUrl + +from fastmcp.resources import FileResource + + +@pytest.fixture +def temp_file(): + """Create a temporary file for testing. + + File is automatically cleaned up after the test if it still exists. + """ + content = "test content" + with NamedTemporaryFile(mode="w", delete=False) as f: + f.write(content) + path = Path(f.name).resolve() + yield path + try: + path.unlink() + except FileNotFoundError: + pass # File was already deleted by the test + + +class TestFileResource: + """Test FileResource functionality.""" + + def test_file_resource_creation(self, temp_file: Path): + """Test creating a FileResource.""" + resource = FileResource( + uri=FileUrl(temp_file.as_uri()), + name="test", + description="test file", + path=temp_file, + ) + assert str(resource.uri) == temp_file.as_uri() + assert resource.name == "test" + assert resource.description == "test file" + assert resource.mime_type == "text/plain" # default + assert resource.path == temp_file + assert resource.is_binary is False # default + + def test_file_resource_str_path_conversion(self, temp_file: Path): + """Test FileResource handles string paths.""" + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=Path(str(temp_file)), + ) + assert isinstance(resource.path, Path) + assert resource.path.is_absolute() + + @pytest.mark.anyio + async def test_read_text_file(self, temp_file: Path): + """Test reading a text file.""" + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + ) + content = await resource.read() + assert content == "test content" + assert resource.mime_type == "text/plain" + + @pytest.mark.anyio + async def test_read_binary_file(self, temp_file: Path): + """Test reading a file as binary.""" + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + is_binary=True, + ) + content = await resource.read() + assert isinstance(content, bytes) + assert content == b"test content" + + def test_relative_path_error(self): + """Test error on relative path.""" + with pytest.raises(ValueError, match="Path must be absolute"): + FileResource( + uri=FileUrl("file:///test.txt"), + name="test", + path=Path("test.txt"), + ) + + @pytest.mark.anyio + async def test_missing_file_error(self, temp_file: Path): + """Test error when file doesn't exist.""" + # Create path to non-existent file + missing = temp_file.parent / "missing.txt" + resource = FileResource( + uri=FileUrl("file:///missing.txt"), + name="test", + path=missing, + ) + with pytest.raises(ValueError, match="Error reading file"): + await resource.read() + + @pytest.mark.skipif( + os.name == "nt", reason="File permissions behave differently on Windows" + ) + @pytest.mark.anyio + async def test_permission_error(self, temp_file: Path): + """Test reading a file without permissions.""" + temp_file.chmod(0o000) # Remove all permissions + try: + resource = FileResource( + uri=FileUrl(temp_file.as_uri()), + name="test", + path=temp_file, + ) + with pytest.raises(ValueError, match="Error reading file"): + await resource.read() + finally: + temp_file.chmod(0o644) # Restore permissions diff --git a/tests/resources/test_function_resources.py b/tests/resources/test_function_resources.py new file mode 100644 index 000000000..5faba0b88 --- /dev/null +++ b/tests/resources/test_function_resources.py @@ -0,0 +1,138 @@ +import pytest +from pydantic import AnyUrl, BaseModel + +from fastmcp.resources import FunctionResource + + +class TestFunctionResource: + """Test FunctionResource functionality.""" + + def test_function_resource_creation(self): + """Test creating a FunctionResource.""" + + def my_func() -> str: + return "test content" + + resource = FunctionResource( + uri=AnyUrl("fn://test"), + name="test", + description="test function", + fn=my_func, + ) + assert str(resource.uri) == "fn://test" + assert resource.name == "test" + assert resource.description == "test function" + assert resource.mime_type == "text/plain" # default + assert resource.fn == my_func + + @pytest.mark.anyio + async def test_read_text(self): + """Test reading text from a FunctionResource.""" + + def get_data() -> str: + return "Hello, world!" + + resource = FunctionResource( + uri=AnyUrl("function://test"), + name="test", + fn=get_data, + ) + content = await resource.read() + assert content == "Hello, world!" + assert resource.mime_type == "text/plain" + + @pytest.mark.anyio + async def test_read_binary(self): + """Test reading binary data from a FunctionResource.""" + + def get_data() -> bytes: + return b"Hello, world!" + + resource = FunctionResource( + uri=AnyUrl("function://test"), + name="test", + fn=get_data, + ) + content = await resource.read() + assert content == b"Hello, world!" + + @pytest.mark.anyio + async def test_json_conversion(self): + """Test automatic JSON conversion of non-string results.""" + + def get_data() -> dict: + return {"key": "value"} + + resource = FunctionResource( + uri=AnyUrl("function://test"), + name="test", + fn=get_data, + ) + content = await resource.read() + assert isinstance(content, str) + assert '"key": "value"' in content + + @pytest.mark.anyio + async def test_error_handling(self): + """Test error handling in FunctionResource.""" + + def failing_func() -> str: + raise ValueError("Test error") + + resource = FunctionResource( + uri=AnyUrl("function://test"), + name="test", + fn=failing_func, + ) + with pytest.raises(ValueError, match="Error reading resource function://test"): + await resource.read() + + @pytest.mark.anyio + async def test_basemodel_conversion(self): + """Test handling of BaseModel types.""" + + class MyModel(BaseModel): + name: str + + resource = FunctionResource( + uri=AnyUrl("function://test"), + name="test", + fn=lambda: MyModel(name="test"), + ) + content = await resource.read() + assert content == '{"name": "test"}' + + @pytest.mark.anyio + async def test_custom_type_conversion(self): + """Test handling of custom types.""" + + class CustomData: + def __str__(self) -> str: + return "custom data" + + def get_data() -> CustomData: + return CustomData() + + resource = FunctionResource( + uri=AnyUrl("function://test"), + name="test", + fn=get_data, + ) + content = await resource.read() + assert isinstance(content, str) + + @pytest.mark.anyio + async def test_async_read_text(self): + """Test reading text from async FunctionResource.""" + + async def get_data() -> str: + return "Hello, world!" + + resource = FunctionResource( + uri=AnyUrl("function://test"), + name="test", + fn=get_data, + ) + content = await resource.read() + assert content == "Hello, world!" + assert resource.mime_type == "text/plain" diff --git a/tests/resources/test_resource_manager.py b/tests/resources/test_resource_manager.py index 6484415e2..12f7cb59d 100644 --- a/tests/resources/test_resource_manager.py +++ b/tests/resources/test_resource_manager.py @@ -1,221 +1,363 @@ -from mcp.server.fastmcp.resources import FunctionResource, ResourceTemplate -from pydantic.networks import AnyUrl +from pathlib import Path +from tempfile import NamedTemporaryFile -from fastmcp.resources.resource_manager import ResourceManager +import pytest +from pydantic import AnyUrl, FileUrl + +from fastmcp.resources import ( + FileResource, + FunctionResource, + ResourceManager, + ResourceTemplate, +) -def test_import_resources(): - """Test importing resources from one manager to another with a prefix.""" - # Setup source manager with resources - source_manager = ResourceManager() +@pytest.fixture +def temp_file(): + """Create a temporary file for testing. - # Create mock resource functions - async def weather_fn(): - return "Weather data" - - async def traffic_fn(): - return "Traffic data" - - # Add resources to source manager - weather_resource = FunctionResource( - uri=AnyUrl("weather://forecast"), - name="weather_forecast", - description="Get weather forecast", - mime_type="application/json", - fn=weather_fn, - ) - source_manager._resources["weather://forecast"] = weather_resource - - traffic_resource = FunctionResource( - uri=AnyUrl("traffic://status"), - name="traffic_status", - description="Get traffic status", - mime_type="application/json", - fn=traffic_fn, - ) - source_manager._resources["traffic://status"] = traffic_resource - - # Create target manager - target_manager = ResourceManager() - - # Import resources from source to target - prefix = "data+" - target_manager.import_resources(source_manager, prefix) - - # Verify resources were imported with prefixes - assert "data+weather://forecast" in target_manager._resources - assert "data+traffic://status" in target_manager._resources - - # Verify the original resources still exist in source manager - assert "weather://forecast" in source_manager._resources - assert "traffic://status" in source_manager._resources - - # Verify the imported resources have the correct properties - assert ( - target_manager._resources["data+weather://forecast"].name == "weather_forecast" - ) - assert ( - target_manager._resources["data+weather://forecast"].description - == "Get weather forecast" - ) - assert ( - target_manager._resources["data+weather://forecast"].mime_type - == "application/json" - ) - - assert target_manager._resources["data+traffic://status"].name == "traffic_status" - assert ( - target_manager._resources["data+traffic://status"].description - == "Get traffic status" - ) - assert ( - target_manager._resources["data+traffic://status"].mime_type - == "application/json" - ) - - # Since we're dealing with FunctionResource type, we can safely check function attributes - assert isinstance( - target_manager._resources["data+weather://forecast"], FunctionResource - ) - assert isinstance( - target_manager._resources["data+traffic://status"], FunctionResource - ) - - weather_resource = target_manager._resources["data+weather://forecast"] - traffic_resource = target_manager._resources["data+traffic://status"] - - if hasattr(weather_resource, "fn") and hasattr(traffic_resource, "fn"): - assert weather_resource.fn.__name__ == weather_fn.__name__ - assert traffic_resource.fn.__name__ == traffic_fn.__name__ + File is automatically cleaned up after the test if it still exists. + """ + content = "test content" + with NamedTemporaryFile(mode="w", delete=False) as f: + f.write(content) + path = Path(f.name).resolve() + yield path + try: + path.unlink() + except FileNotFoundError: + pass # File was already deleted by the test -def test_import_templates(): - """Test importing resource templates from one manager to another with a prefix.""" - # Setup source manager with templates - source_manager = ResourceManager() +class TestResourceManager: + """Test ResourceManager functionality.""" - # Create mock template functions - async def user_fn(**params): - return f"User data for id {params.get('id')}" + def test_add_resource(self, temp_file: Path): + """Test adding a resource.""" + manager = ResourceManager() + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + ) + added = manager.add_resource(resource) + assert added == resource + assert manager.list_resources() == [resource] - async def product_fn(**params): - return f"Product data for id {params.get('id')}" + def test_add_duplicate_resource(self, temp_file: Path): + """Test adding the same resource twice.""" + manager = ResourceManager() + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + ) + first = manager.add_resource(resource) + second = manager.add_resource(resource) + assert first == second + assert manager.list_resources() == [resource] - # Add templates to source manager - user_template = ResourceTemplate( - uri_template="api://users/{id}", - name="user_template", - description="Get user by ID", - mime_type="application/json", - fn=user_fn, - parameters={"id": {"type": "string", "description": "User ID"}}, - ) - source_manager._templates["api://users/{id}"] = user_template + def test_warn_on_duplicate_resources(self, temp_file: Path, caplog): + """Test warning on duplicate resources.""" + manager = ResourceManager() + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + ) + manager.add_resource(resource) + manager.add_resource(resource) + assert "Resource already exists" in caplog.text - product_template = ResourceTemplate( - uri_template="api://products/{id}", - name="product_template", - description="Get product by ID", - mime_type="application/json", - fn=product_fn, - parameters={"id": {"type": "string", "description": "Product ID"}}, - ) - source_manager._templates["api://products/{id}"] = product_template + def test_disable_warn_on_duplicate_resources(self, temp_file: Path, caplog): + """Test disabling warning on duplicate resources.""" + manager = ResourceManager(warn_on_duplicate_resources=False) + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + ) + manager.add_resource(resource) + manager.add_resource(resource) + assert "Resource already exists" not in caplog.text - # Create target manager - target_manager = ResourceManager() + @pytest.mark.anyio + async def test_get_resource(self, temp_file: Path): + """Test getting a resource by URI.""" + manager = ResourceManager() + resource = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test", + path=temp_file, + ) + manager.add_resource(resource) + retrieved = await manager.get_resource(resource.uri) + assert retrieved == resource - # Import templates from source to target - prefix = "shop+" - target_manager.import_templates(source_manager, prefix) + @pytest.mark.anyio + async def test_get_resource_from_template(self): + """Test getting a resource through a template.""" + manager = ResourceManager() - # Verify templates were imported with prefixes - assert "shop+api://users/{id}" in target_manager._templates - assert "shop+api://products/{id}" in target_manager._templates + def greet(name: str) -> str: + return f"Hello, {name}!" - # Verify the original templates still exist in source manager - assert "api://users/{id}" in source_manager._templates - assert "api://products/{id}" in source_manager._templates + template = ResourceTemplate.from_function( + fn=greet, + uri_template="greet://{name}", + name="greeter", + ) + manager._templates[template.uri_template] = template - # Verify the imported templates have the correct properties - assert target_manager._templates["shop+api://users/{id}"].name == "user_template" - assert ( - target_manager._templates["shop+api://users/{id}"].description - == "Get user by ID" - ) - assert ( - target_manager._templates["shop+api://users/{id}"].mime_type - == "application/json" - ) - assert target_manager._templates["shop+api://users/{id}"].parameters == { - "id": {"type": "string", "description": "User ID"} - } + resource = await manager.get_resource(AnyUrl("greet://world")) + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Hello, world!" - assert ( - target_manager._templates["shop+api://products/{id}"].name == "product_template" - ) - assert ( - target_manager._templates["shop+api://products/{id}"].description - == "Get product by ID" - ) - assert ( - target_manager._templates["shop+api://products/{id}"].mime_type - == "application/json" - ) - assert target_manager._templates["shop+api://products/{id}"].parameters == { - "id": {"type": "string", "description": "Product ID"} - } + @pytest.mark.anyio + async def test_get_unknown_resource(self): + """Test getting a non-existent resource.""" + manager = ResourceManager() + with pytest.raises(ValueError, match="Unknown resource"): + await manager.get_resource(AnyUrl("unknown://test")) - # Verify the template functions were properly copied (only if the fn attribute exists) - user_template = target_manager._templates["shop+api://users/{id}"] - product_template = target_manager._templates["shop+api://products/{id}"] - - if hasattr(user_template, "fn") and hasattr(product_template, "fn"): - assert user_template.fn.__name__ == user_fn.__name__ - assert product_template.fn.__name__ == product_fn.__name__ + def test_list_resources(self, temp_file: Path): + """Test listing all resources.""" + manager = ResourceManager() + resource1 = FileResource( + uri=FileUrl(f"file://{temp_file}"), + name="test1", + path=temp_file, + ) + resource2 = FileResource( + uri=FileUrl(f"file://{temp_file}2"), + name="test2", + path=temp_file, + ) + manager.add_resource(resource1) + manager.add_resource(resource2) + resources = manager.list_resources() + assert len(resources) == 2 + assert resources == [resource1, resource2] -def test_import_multiple_resource_types(): - """Test importing both resources and templates with the same prefix.""" - # Setup source manager with both resources and templates - source_manager = ResourceManager() +class TestImports: + def test_import_resources(self): + """Test importing resources from one manager to another with a prefix.""" + # Setup source manager with resources + source_manager = ResourceManager() - # Create mock functions - async def resource_fn(): - return "Resource data" + # Create mock resource functions + async def weather_fn(): + return "Weather data" - async def template_fn(**params): - return f"Template data for id {params.get('id')}" + async def traffic_fn(): + return "Traffic data" - # Add a resource to source manager - resource = FunctionResource( - uri=AnyUrl("data://resource"), - name="test_resource", - description="Test resource", - mime_type="application/json", - fn=resource_fn, - ) - source_manager._resources["data://resource"] = resource + # Add resources to source manager + weather_resource = FunctionResource( + uri=AnyUrl("weather://forecast"), + name="weather_forecast", + description="Get weather forecast", + mime_type="application/json", + fn=weather_fn, + ) + source_manager._resources["weather://forecast"] = weather_resource - # Add a template to source manager - template = ResourceTemplate( - uri_template="data://template/{id}", - name="test_template", - description="Test template", - mime_type="application/json", - fn=template_fn, - parameters={"id": {"type": "string", "description": "ID parameter"}}, - ) - source_manager._templates["data://template/{id}"] = template + traffic_resource = FunctionResource( + uri=AnyUrl("traffic://status"), + name="traffic_status", + description="Get traffic status", + mime_type="application/json", + fn=traffic_fn, + ) + source_manager._resources["traffic://status"] = traffic_resource - # Create target manager - target_manager = ResourceManager() + # Create target manager + target_manager = ResourceManager() - # Import both resources and templates - prefix = "test+" - target_manager.import_resources(source_manager, prefix) - target_manager.import_templates(source_manager, prefix) + # Import resources from source to target + prefix = "data+" + target_manager.import_resources(source_manager, prefix) - # Verify both resource types were imported with prefixes - assert "test+data://resource" in target_manager._resources - assert "test+data://template/{id}" in target_manager._templates + # Verify resources were imported with prefixes + assert "data+weather://forecast" in target_manager._resources + assert "data+traffic://status" in target_manager._resources + + # Verify the original resources still exist in source manager + assert "weather://forecast" in source_manager._resources + assert "traffic://status" in source_manager._resources + + # Verify the imported resources have the correct properties + assert ( + target_manager._resources["data+weather://forecast"].name + == "weather_forecast" + ) + assert ( + target_manager._resources["data+weather://forecast"].description + == "Get weather forecast" + ) + assert ( + target_manager._resources["data+weather://forecast"].mime_type + == "application/json" + ) + + assert ( + target_manager._resources["data+traffic://status"].name == "traffic_status" + ) + assert ( + target_manager._resources["data+traffic://status"].description + == "Get traffic status" + ) + assert ( + target_manager._resources["data+traffic://status"].mime_type + == "application/json" + ) + + # Since we're dealing with FunctionResource type, we can safely check function attributes + assert isinstance( + target_manager._resources["data+weather://forecast"], FunctionResource + ) + assert isinstance( + target_manager._resources["data+traffic://status"], FunctionResource + ) + + weather_resource = target_manager._resources["data+weather://forecast"] + traffic_resource = target_manager._resources["data+traffic://status"] + + if hasattr(weather_resource, "fn") and hasattr(traffic_resource, "fn"): + assert weather_resource.fn.__name__ == weather_fn.__name__ + assert traffic_resource.fn.__name__ == traffic_fn.__name__ + + def test_import_templates(self): + """Test importing resource templates from one manager to another with a prefix.""" + # Setup source manager with templates + source_manager = ResourceManager() + + # Create mock template functions + async def user_fn(**params): + return f"User data for id {params.get('id')}" + + async def product_fn(**params): + return f"Product data for id {params.get('id')}" + + # Add templates to source manager + user_template = ResourceTemplate( + uri_template="api://users/{id}", + name="user_template", + description="Get user by ID", + mime_type="application/json", + fn=user_fn, + parameters={"id": {"type": "string", "description": "User ID"}}, + ) + source_manager._templates["api://users/{id}"] = user_template + + product_template = ResourceTemplate( + uri_template="api://products/{id}", + name="product_template", + description="Get product by ID", + mime_type="application/json", + fn=product_fn, + parameters={"id": {"type": "string", "description": "Product ID"}}, + ) + source_manager._templates["api://products/{id}"] = product_template + + # Create target manager + target_manager = ResourceManager() + + # Import templates from source to target + prefix = "shop+" + target_manager.import_templates(source_manager, prefix) + + # Verify templates were imported with prefixes + assert "shop+api://users/{id}" in target_manager._templates + assert "shop+api://products/{id}" in target_manager._templates + + # Verify the original templates still exist in source manager + assert "api://users/{id}" in source_manager._templates + assert "api://products/{id}" in source_manager._templates + + # Verify the imported templates have the correct properties + assert ( + target_manager._templates["shop+api://users/{id}"].name == "user_template" + ) + assert ( + target_manager._templates["shop+api://users/{id}"].description + == "Get user by ID" + ) + assert ( + target_manager._templates["shop+api://users/{id}"].mime_type + == "application/json" + ) + assert target_manager._templates["shop+api://users/{id}"].parameters == { + "id": {"type": "string", "description": "User ID"} + } + + assert ( + target_manager._templates["shop+api://products/{id}"].name + == "product_template" + ) + assert ( + target_manager._templates["shop+api://products/{id}"].description + == "Get product by ID" + ) + assert ( + target_manager._templates["shop+api://products/{id}"].mime_type + == "application/json" + ) + assert target_manager._templates["shop+api://products/{id}"].parameters == { + "id": {"type": "string", "description": "Product ID"} + } + + # Verify the template functions were properly copied (only if the fn attribute exists) + user_template = target_manager._templates["shop+api://users/{id}"] + product_template = target_manager._templates["shop+api://products/{id}"] + + if hasattr(user_template, "fn") and hasattr(product_template, "fn"): + assert user_template.fn.__name__ == user_fn.__name__ + assert product_template.fn.__name__ == product_fn.__name__ + + def test_import_multiple_resource_types(self): + """Test importing both resources and templates with the same prefix.""" + # Setup source manager with both resources and templates + source_manager = ResourceManager() + + # Create mock functions + async def resource_fn(): + return "Resource data" + + async def template_fn(**params): + return f"Template data for id {params.get('id')}" + + # Add a resource to source manager + resource = FunctionResource( + uri=AnyUrl("data://resource"), + name="test_resource", + description="Test resource", + mime_type="application/json", + fn=resource_fn, + ) + source_manager._resources["data://resource"] = resource + + # Add a template to source manager + template = ResourceTemplate( + uri_template="data://template/{id}", + name="test_template", + description="Test template", + mime_type="application/json", + fn=template_fn, + parameters={"id": {"type": "string", "description": "ID parameter"}}, + ) + source_manager._templates["data://template/{id}"] = template + + # Create target manager + target_manager = ResourceManager() + + # Import both resources and templates + prefix = "test+" + target_manager.import_resources(source_manager, prefix) + target_manager.import_templates(source_manager, prefix) + + # Verify both resource types were imported with prefixes + assert "test+data://resource" in target_manager._resources + assert "test+data://template/{id}" in target_manager._templates diff --git a/tests/resources/test_resource_template.py b/tests/resources/test_resource_template.py new file mode 100644 index 000000000..2d9d9b78f --- /dev/null +++ b/tests/resources/test_resource_template.py @@ -0,0 +1,188 @@ +import json + +import pytest +from pydantic import BaseModel + +from fastmcp.resources import FunctionResource, ResourceTemplate + + +class TestResourceTemplate: + """Test ResourceTemplate functionality.""" + + def test_template_creation(self): + """Test creating a template from a function.""" + + def my_func(key: str, value: int) -> dict: + return {"key": key, "value": value} + + template = ResourceTemplate.from_function( + fn=my_func, + uri_template="test://{key}/{value}", + name="test", + ) + assert template.uri_template == "test://{key}/{value}" + assert template.name == "test" + assert template.mime_type == "text/plain" # default + test_input = {"key": "test", "value": 42} + assert template.fn(**test_input) == my_func(**test_input) + + def test_template_matches(self): + """Test matching URIs against a template.""" + + def my_func(key: str, value: int) -> dict: + return {"key": key, "value": value} + + template = ResourceTemplate.from_function( + fn=my_func, + uri_template="test://{key}/{value}", + name="test", + ) + + # Valid match + params = template.matches("test://foo/123") + assert params == {"key": "foo", "value": "123"} + + # No match + assert template.matches("test://foo") is None + assert template.matches("other://foo/123") is None + + @pytest.mark.anyio + async def test_create_resource(self): + """Test creating a resource from a template.""" + + def my_func(key: str, value: int) -> dict: + return {"key": key, "value": value} + + template = ResourceTemplate.from_function( + fn=my_func, + uri_template="test://{key}/{value}", + name="test", + ) + + resource = await template.create_resource( + "test://foo/123", + {"key": "foo", "value": 123}, + ) + + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert isinstance(content, str) + data = json.loads(content) + assert data == {"key": "foo", "value": 123} + + @pytest.mark.anyio + async def test_template_error(self): + """Test error handling in template resource creation.""" + + def failing_func(x: str) -> str: + raise ValueError("Test error") + + template = ResourceTemplate.from_function( + fn=failing_func, + uri_template="fail://{x}", + name="fail", + ) + + with pytest.raises(ValueError, match="Error creating resource from template"): + await template.create_resource("fail://test", {"x": "test"}) + + @pytest.mark.anyio + async def test_async_text_resource(self): + """Test creating a text resource from async function.""" + + async def greet(name: str) -> str: + return f"Hello, {name}!" + + template = ResourceTemplate.from_function( + fn=greet, + uri_template="greet://{name}", + name="greeter", + ) + + resource = await template.create_resource( + "greet://world", + {"name": "world"}, + ) + + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "Hello, world!" + + @pytest.mark.anyio + async def test_async_binary_resource(self): + """Test creating a binary resource from async function.""" + + async def get_bytes(value: str) -> bytes: + return value.encode() + + template = ResourceTemplate.from_function( + fn=get_bytes, + uri_template="bytes://{value}", + name="bytes", + ) + + resource = await template.create_resource( + "bytes://test", + {"value": "test"}, + ) + + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == b"test" + + @pytest.mark.anyio + async def test_basemodel_conversion(self): + """Test handling of BaseModel types.""" + + class MyModel(BaseModel): + key: str + value: int + + def get_data(key: str, value: int) -> MyModel: + return MyModel(key=key, value=value) + + template = ResourceTemplate.from_function( + fn=get_data, + uri_template="test://{key}/{value}", + name="test", + ) + + resource = await template.create_resource( + "test://foo/123", + {"key": "foo", "value": 123}, + ) + + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert isinstance(content, str) + data = json.loads(content) + assert data == {"key": "foo", "value": 123} + + @pytest.mark.anyio + async def test_custom_type_conversion(self): + """Test handling of custom types.""" + + class CustomData: + def __init__(self, value: str): + self.value = value + + def __str__(self) -> str: + return self.value + + def get_data(value: str) -> CustomData: + return CustomData(value) + + template = ResourceTemplate.from_function( + fn=get_data, + uri_template="test://{value}", + name="test", + ) + + resource = await template.create_resource( + "test://hello", + {"value": "hello"}, + ) + + assert isinstance(resource, FunctionResource) + content = await resource.read() + assert content == "hello" diff --git a/tests/resources/test_resources.py b/tests/resources/test_resources.py new file mode 100644 index 000000000..870002c35 --- /dev/null +++ b/tests/resources/test_resources.py @@ -0,0 +1,101 @@ +import pytest +from pydantic import AnyUrl + +from fastmcp.resources import FunctionResource, Resource + + +class TestResourceValidation: + """Test base Resource validation.""" + + def test_resource_uri_validation(self): + """Test URI validation.""" + + def dummy_func() -> str: + return "data" + + # Valid URI + resource = FunctionResource( + uri=AnyUrl("http://example.com/data"), + name="test", + fn=dummy_func, + ) + assert str(resource.uri) == "http://example.com/data" + + # Missing protocol + with pytest.raises(ValueError, match="Input should be a valid URL"): + FunctionResource( + uri=AnyUrl("invalid"), + name="test", + fn=dummy_func, + ) + + # Missing host + with pytest.raises(ValueError, match="Input should be a valid URL"): + FunctionResource( + uri=AnyUrl("http://"), + name="test", + fn=dummy_func, + ) + + def test_resource_name_from_uri(self): + """Test name is extracted from URI if not provided.""" + + def dummy_func() -> str: + return "data" + + resource = FunctionResource( + uri=AnyUrl("resource://my-resource"), + fn=dummy_func, + ) + assert resource.name == "resource://my-resource" + + def test_resource_name_validation(self): + """Test name validation.""" + + def dummy_func() -> str: + return "data" + + # Must provide either name or URI + with pytest.raises(ValueError, match="Either name or uri must be provided"): + FunctionResource( + fn=dummy_func, + ) + + # Explicit name takes precedence over URI + resource = FunctionResource( + uri=AnyUrl("resource://uri-name"), + name="explicit-name", + fn=dummy_func, + ) + assert resource.name == "explicit-name" + + def test_resource_mime_type(self): + """Test mime type handling.""" + + def dummy_func() -> str: + return "data" + + # Default mime type + resource = FunctionResource( + uri=AnyUrl("resource://test"), + fn=dummy_func, + ) + assert resource.mime_type == "text/plain" + + # Custom mime type + resource = FunctionResource( + uri=AnyUrl("resource://test"), + fn=dummy_func, + mime_type="application/json", + ) + assert resource.mime_type == "application/json" + + @pytest.mark.anyio + async def test_resource_read_abstract(self): + """Test that Resource.read() is abstract.""" + + class ConcreteResource(Resource): + pass + + with pytest.raises(TypeError, match="abstract method"): + ConcreteResource(uri=AnyUrl("test://test"), name="test") # type: ignore diff --git a/tests/server/__init__.py b/tests/server/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/test_file_server.py b/tests/server/test_file_server.py new file mode 100644 index 000000000..b55ea8533 --- /dev/null +++ b/tests/server/test_file_server.py @@ -0,0 +1,132 @@ +import json +from pathlib import Path + +import pytest + +from fastmcp import FastMCP + + +@pytest.fixture() +def test_dir(tmp_path_factory) -> Path: + """Create a temporary directory with test files.""" + tmp = tmp_path_factory.mktemp("test_files") + + # Create test files + (tmp / "example.py").write_text("print('hello world')") + (tmp / "readme.md").write_text("# Test Directory\nThis is a test.") + (tmp / "config.json").write_text('{"test": true}') + + return tmp + + +@pytest.fixture +def mcp() -> FastMCP: + mcp = FastMCP() + + return mcp + + +@pytest.fixture(autouse=True) +def resources(mcp: FastMCP, test_dir: Path) -> FastMCP: + @mcp.resource("dir://test_dir") + def list_test_dir() -> list[str]: + """List the files in the test directory""" + return [str(f) for f in test_dir.iterdir()] + + @mcp.resource("file://test_dir/example.py") + def read_example_py() -> str: + """Read the example.py file""" + try: + return (test_dir / "example.py").read_text() + except FileNotFoundError: + return "File not found" + + @mcp.resource("file://test_dir/readme.md") + def read_readme_md() -> str: + """Read the readme.md file""" + try: + return (test_dir / "readme.md").read_text() + except FileNotFoundError: + return "File not found" + + @mcp.resource("file://test_dir/config.json") + def read_config_json() -> str: + """Read the config.json file""" + try: + return (test_dir / "config.json").read_text() + except FileNotFoundError: + return "File not found" + + return mcp + + +@pytest.fixture(autouse=True) +def tools(mcp: FastMCP, test_dir: Path) -> FastMCP: + @mcp.tool() + def delete_file(path: str) -> bool: + # ensure path is in test_dir + if Path(path).resolve().parent != test_dir: + raise ValueError(f"Path must be in test_dir: {path}") + Path(path).unlink() + return True + + return mcp + + +@pytest.mark.anyio +async def test_list_resources(mcp: FastMCP): + resources = await mcp.list_resources() + assert len(resources) == 4 + + assert [str(r.uri) for r in resources] == [ + "dir://test_dir", + "file://test_dir/example.py", + "file://test_dir/readme.md", + "file://test_dir/config.json", + ] + + +@pytest.mark.anyio +async def test_read_resource_dir(mcp: FastMCP): + res_iter = await mcp.read_resource("dir://test_dir") + res_list = list(res_iter) + assert len(res_list) == 1 + res = res_list[0] + assert res.mime_type == "text/plain" + + files = json.loads(res.content) + + assert sorted([Path(f).name for f in files]) == [ + "config.json", + "example.py", + "readme.md", + ] + + +@pytest.mark.anyio +async def test_read_resource_file(mcp: FastMCP): + res_iter = await mcp.read_resource("file://test_dir/example.py") + res_list = list(res_iter) + assert len(res_list) == 1 + res = res_list[0] + assert res.content == "print('hello world')" + + +@pytest.mark.anyio +async def test_delete_file(mcp: FastMCP, test_dir: Path): + await mcp.call_tool( + "delete_file", arguments=dict(path=str(test_dir / "example.py")) + ) + assert not (test_dir / "example.py").exists() + + +@pytest.mark.anyio +async def test_delete_file_and_check_resources(mcp: FastMCP, test_dir: Path): + await mcp.call_tool( + "delete_file", arguments=dict(path=str(test_dir / "example.py")) + ) + res_iter = await mcp.read_resource("file://test_dir/example.py") + res_list = list(res_iter) + assert len(res_list) == 1 + res = res_list[0] + assert res.content == "File not found" diff --git a/tests/server/test_lifespan.py b/tests/server/test_lifespan.py new file mode 100644 index 000000000..30598d4ad --- /dev/null +++ b/tests/server/test_lifespan.py @@ -0,0 +1,114 @@ +"""Tests for lifespan functionality in both low-level and FastMCP servers.""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +import anyio +import pytest +from mcp.types import ( + ClientCapabilities, + Implementation, + InitializeRequestParams, + JSONRPCMessage, + JSONRPCNotification, + JSONRPCRequest, +) +from pydantic import TypeAdapter + +from fastmcp import Context, FastMCP + + +@pytest.mark.anyio +async def test_fastmcp_server_lifespan(): + """Test that lifespan works in FastMCP server.""" + + @asynccontextmanager + async def test_lifespan(server: FastMCP) -> AsyncIterator[dict]: + """Test lifespan context that tracks startup/shutdown.""" + context = {"started": False, "shutdown": False} + try: + context["started"] = True + yield context + finally: + context["shutdown"] = True + + server = FastMCP("test", lifespan=test_lifespan) + + # Create memory streams for testing + send_stream1, receive_stream1 = anyio.create_memory_object_stream(100) + send_stream2, receive_stream2 = anyio.create_memory_object_stream(100) + + # Add a tool that checks lifespan context + @server.tool() + def check_lifespan(ctx: Context) -> bool: + """Tool that checks lifespan context.""" + assert isinstance(ctx.request_context.lifespan_context, dict) + assert ctx.request_context.lifespan_context["started"] + assert not ctx.request_context.lifespan_context["shutdown"] + return True + + # Run server in background task + async with ( + anyio.create_task_group() as tg, + send_stream1, + receive_stream1, + send_stream2, + receive_stream2, + ): + + async def run_server(): + await server._mcp_server.run( + receive_stream1, + send_stream2, + server._mcp_server.create_initialization_options(), + raise_exceptions=True, + ) + + tg.start_soon(run_server) + + # Initialize the server + params = InitializeRequestParams( + protocolVersion="2024-11-05", + capabilities=ClientCapabilities(), + clientInfo=Implementation(name="test-client", version="0.1.0"), + ) + await send_stream1.send( + JSONRPCMessage( + root=JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="initialize", + params=TypeAdapter(InitializeRequestParams).dump_python(params), + ) + ) + ) + response = await receive_stream2.receive() + + # Send initialized notification + await send_stream1.send( + JSONRPCMessage( + root=JSONRPCNotification( + jsonrpc="2.0", + method="notifications/initialized", + ) + ) + ) + + # Call the tool to verify lifespan context + await send_stream1.send( + JSONRPCMessage( + root=JSONRPCRequest( + jsonrpc="2.0", + id=2, + method="tools/call", + params={"name": "check_lifespan", "arguments": {}}, + ) + ) + ) + + # Get response and verify + response = await receive_stream2.receive() + assert response.root.result["content"][0]["text"] == "true" + + # Cancel server task + tg.cancel_scope.cancel() diff --git a/tests/server/test_server.py b/tests/server/test_server.py new file mode 100644 index 000000000..d79ce8f6d --- /dev/null +++ b/tests/server/test_server.py @@ -0,0 +1,770 @@ +import base64 +from pathlib import Path +from typing import TYPE_CHECKING + +import pytest +from mcp.shared.exceptions import McpError +from mcp.shared.memory import ( + create_connected_server_and_client_session as client_session, +) +from mcp.types import ( + BlobResourceContents, + ImageContent, + TextContent, + TextResourceContents, +) +from pydantic import AnyUrl, Field + +from fastmcp import Context, FastMCP +from fastmcp.prompts.base import EmbeddedResource, Message, UserMessage +from fastmcp.resources import FileResource, FunctionResource +from fastmcp.utilities.types import Image + +if TYPE_CHECKING: + from fastmcp import Context + + +class TestServer: + @pytest.mark.anyio + async def test_create_server(self): + mcp = FastMCP(instructions="Server instructions") + assert mcp.name == "FastMCP" + assert mcp.instructions == "Server instructions" + + @pytest.mark.anyio + async def test_non_ascii_description(self): + """Test that FastMCP handles non-ASCII characters in descriptions correctly""" + mcp = FastMCP() + + @mcp.tool( + description=( + "🌟 This tool uses emojis and UTF-8 characters: á é í ó ú ñ 漢字 🎉" + ) + ) + def hello_world(name: str = "世界") -> str: + return f"¡Hola, {name}! 👋" + + async with client_session(mcp._mcp_server) as client: + tools = await client.list_tools() + assert len(tools.tools) == 1 + tool = tools.tools[0] + assert tool.description is not None + assert "🌟" in tool.description + assert "漢字" in tool.description + assert "🎉" in tool.description + + result = await client.call_tool("hello_world", {}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert "¡Hola, 世界! 👋" == content.text + + @pytest.mark.anyio + async def test_add_tool_decorator(self): + mcp = FastMCP() + + @mcp.tool() + def add(x: int, y: int) -> int: + return x + y + + assert len(mcp._tool_manager.list_tools()) == 1 + + @pytest.mark.anyio + async def test_add_tool_decorator_incorrect_usage(self): + mcp = FastMCP() + + with pytest.raises(TypeError, match="The @tool decorator was used incorrectly"): + + @mcp.tool # Missing parentheses #type: ignore + def add(x: int, y: int) -> int: + return x + y + + @pytest.mark.anyio + async def test_add_resource_decorator(self): + mcp = FastMCP() + + @mcp.resource("r://{x}") + def get_data(x: str) -> str: + return f"Data: {x}" + + assert len(mcp._resource_manager._templates) == 1 + + @pytest.mark.anyio + async def test_add_resource_decorator_incorrect_usage(self): + mcp = FastMCP() + + with pytest.raises( + TypeError, match="The @resource decorator was used incorrectly" + ): + + @mcp.resource # Missing parentheses #type: ignore + def get_data(x: str) -> str: + return f"Data: {x}" + + +def tool_fn(x: int, y: int) -> int: + return x + y + + +def error_tool_fn() -> None: + raise ValueError("Test error") + + +def image_tool_fn(path: str) -> Image: + return Image(path) + + +def mixed_content_tool_fn() -> list[TextContent | ImageContent]: + return [ + TextContent(type="text", text="Hello"), + ImageContent(type="image", data="abc", mimeType="image/png"), + ] + + +class TestServerTools: + @pytest.mark.anyio + async def test_add_tool(self): + mcp = FastMCP() + mcp.add_tool(tool_fn) + mcp.add_tool(tool_fn) + assert len(mcp._tool_manager.list_tools()) == 1 + + @pytest.mark.anyio + async def test_list_tools(self): + mcp = FastMCP() + mcp.add_tool(tool_fn) + async with client_session(mcp._mcp_server) as client: + tools = await client.list_tools() + assert len(tools.tools) == 1 + + @pytest.mark.anyio + async def test_call_tool(self): + mcp = FastMCP() + mcp.add_tool(tool_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("my_tool", {"arg1": "value"}) + assert not hasattr(result, "error") + assert len(result.content) > 0 + + @pytest.mark.anyio + async def test_tool_exception_handling(self): + mcp = FastMCP() + mcp.add_tool(error_tool_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("error_tool_fn", {}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert "Test error" in content.text + assert result.isError is True + + @pytest.mark.anyio + async def test_tool_error_handling(self): + mcp = FastMCP() + mcp.add_tool(error_tool_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("error_tool_fn", {}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert "Test error" in content.text + assert result.isError is True + + @pytest.mark.anyio + async def test_tool_error_details(self): + """Test that exception details are properly formatted in the response""" + mcp = FastMCP() + mcp.add_tool(error_tool_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("error_tool_fn", {}) + content = result.content[0] + assert isinstance(content, TextContent) + assert isinstance(content.text, str) + assert "Test error" in content.text + assert result.isError is True + + @pytest.mark.anyio + async def test_tool_return_value_conversion(self): + mcp = FastMCP() + mcp.add_tool(tool_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("tool_fn", {"x": 1, "y": 2}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert content.text == "3" + + @pytest.mark.anyio + async def test_tool_image_helper(self, tmp_path: Path): + # Create a test image + image_path = tmp_path / "test.png" + image_path.write_bytes(b"fake png data") + + mcp = FastMCP() + mcp.add_tool(image_tool_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("image_tool_fn", {"path": str(image_path)}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, ImageContent) + assert content.type == "image" + assert content.mimeType == "image/png" + # Verify base64 encoding + decoded = base64.b64decode(content.data) + assert decoded == b"fake png data" + + @pytest.mark.anyio + async def test_tool_mixed_content(self): + mcp = FastMCP() + mcp.add_tool(mixed_content_tool_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("mixed_content_tool_fn", {}) + assert len(result.content) == 2 + content1 = result.content[0] + content2 = result.content[1] + assert isinstance(content1, TextContent) + assert content1.text == "Hello" + assert isinstance(content2, ImageContent) + assert content2.mimeType == "image/png" + assert content2.data == "abc" + + @pytest.mark.anyio + async def test_tool_mixed_list_with_image(self, tmp_path: Path): + """Test that lists containing Image objects and other types are handled + correctly""" + # Create a test image + image_path = tmp_path / "test.png" + image_path.write_bytes(b"test image data") + + def mixed_list_fn() -> list: + return [ + "text message", + Image(image_path), + {"key": "value"}, + TextContent(type="text", text="direct content"), + ] + + mcp = FastMCP() + mcp.add_tool(mixed_list_fn) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("mixed_list_fn", {}) + assert len(result.content) == 4 + # Check text conversion + content1 = result.content[0] + assert isinstance(content1, TextContent) + assert content1.text == "text message" + # Check image conversion + content2 = result.content[1] + assert isinstance(content2, ImageContent) + assert content2.mimeType == "image/png" + assert base64.b64decode(content2.data) == b"test image data" + # Check dict conversion + content3 = result.content[2] + assert isinstance(content3, TextContent) + assert '"key": "value"' in content3.text + # Check direct TextContent + content4 = result.content[3] + assert isinstance(content4, TextContent) + assert content4.text == "direct content" + + async def test_parameter_descriptions(self): + mcp = FastMCP("Test Server") + + @mcp.tool() + def greet( + name: str = Field(description="The name to greet"), + title: str = Field(description="Optional title", default=""), + ) -> str: + """A greeting tool""" + return f"Hello {title} {name}" + + tools = await mcp.list_tools() + assert len(tools) == 1 + tool = tools[0] + + # Check that parameter descriptions are present in the schema + properties = tool.inputSchema["properties"] + assert "name" in properties + assert properties["name"]["description"] == "The name to greet" + assert "title" in properties + assert properties["title"]["description"] == "Optional title" + + +class TestServerResources: + @pytest.mark.anyio + async def test_text_resource(self): + mcp = FastMCP() + + def get_text(): + return "Hello, world!" + + resource = FunctionResource( + uri=AnyUrl("resource://test"), name="test", fn=get_text + ) + mcp.add_resource(resource) + + async with client_session(mcp._mcp_server) as client: + result = await client.read_resource(AnyUrl("resource://test")) + assert isinstance(result.contents[0], TextResourceContents) + assert result.contents[0].text == "Hello, world!" + + @pytest.mark.anyio + async def test_binary_resource(self): + mcp = FastMCP() + + def get_binary(): + return b"Binary data" + + resource = FunctionResource( + uri=AnyUrl("resource://binary"), + name="binary", + fn=get_binary, + mime_type="application/octet-stream", + ) + mcp.add_resource(resource) + + async with client_session(mcp._mcp_server) as client: + result = await client.read_resource(AnyUrl("resource://binary")) + assert isinstance(result.contents[0], BlobResourceContents) + assert result.contents[0].blob == base64.b64encode(b"Binary data").decode() + + @pytest.mark.anyio + async def test_file_resource_text(self, tmp_path: Path): + mcp = FastMCP() + + # Create a text file + text_file = tmp_path / "test.txt" + text_file.write_text("Hello from file!") + + resource = FileResource( + uri=AnyUrl("file://test.txt"), name="test.txt", path=text_file + ) + mcp.add_resource(resource) + + async with client_session(mcp._mcp_server) as client: + result = await client.read_resource(AnyUrl("file://test.txt")) + assert isinstance(result.contents[0], TextResourceContents) + assert result.contents[0].text == "Hello from file!" + + @pytest.mark.anyio + async def test_file_resource_binary(self, tmp_path: Path): + mcp = FastMCP() + + # Create a binary file + binary_file = tmp_path / "test.bin" + binary_file.write_bytes(b"Binary file data") + + resource = FileResource( + uri=AnyUrl("file://test.bin"), + name="test.bin", + path=binary_file, + mime_type="application/octet-stream", + ) + mcp.add_resource(resource) + + async with client_session(mcp._mcp_server) as client: + result = await client.read_resource(AnyUrl("file://test.bin")) + assert isinstance(result.contents[0], BlobResourceContents) + assert ( + result.contents[0].blob + == base64.b64encode(b"Binary file data").decode() + ) + + +class TestServerResourceTemplates: + @pytest.mark.anyio + async def test_resource_with_params(self): + """Test that a resource with function parameters raises an error if the URI + parameters don't match""" + mcp = FastMCP() + + with pytest.raises(ValueError, match="Mismatch between URI parameters"): + + @mcp.resource("resource://data") + def get_data_fn(param: str) -> str: + return f"Data: {param}" + + @pytest.mark.anyio + async def test_resource_with_uri_params(self): + """Test that a resource with URI parameters is automatically a template""" + mcp = FastMCP() + + with pytest.raises(ValueError, match="Mismatch between URI parameters"): + + @mcp.resource("resource://{param}") + def get_data() -> str: + return "Data" + + @pytest.mark.anyio + async def test_resource_with_untyped_params(self): + """Test that a resource with untyped parameters raises an error""" + mcp = FastMCP() + + @mcp.resource("resource://{param}") + def get_data(param) -> str: + return "Data" + + @pytest.mark.anyio + async def test_resource_matching_params(self): + """Test that a resource with matching URI and function parameters works""" + mcp = FastMCP() + + @mcp.resource("resource://{name}/data") + def get_data(name: str) -> str: + return f"Data for {name}" + + async with client_session(mcp._mcp_server) as client: + result = await client.read_resource(AnyUrl("resource://test/data")) + assert isinstance(result.contents[0], TextResourceContents) + assert result.contents[0].text == "Data for test" + + @pytest.mark.anyio + async def test_resource_mismatched_params(self): + """Test that mismatched parameters raise an error""" + mcp = FastMCP() + + with pytest.raises(ValueError, match="Mismatch between URI parameters"): + + @mcp.resource("resource://{name}/data") + def get_data(user: str) -> str: + return f"Data for {user}" + + @pytest.mark.anyio + async def test_resource_multiple_params(self): + """Test that multiple parameters work correctly""" + mcp = FastMCP() + + @mcp.resource("resource://{org}/{repo}/data") + def get_data(org: str, repo: str) -> str: + return f"Data for {org}/{repo}" + + async with client_session(mcp._mcp_server) as client: + result = await client.read_resource( + AnyUrl("resource://cursor/fastmcp/data") + ) + assert isinstance(result.contents[0], TextResourceContents) + assert result.contents[0].text == "Data for cursor/fastmcp" + + @pytest.mark.anyio + async def test_resource_multiple_mismatched_params(self): + """Test that mismatched parameters raise an error""" + mcp = FastMCP() + + with pytest.raises(ValueError, match="Mismatch between URI parameters"): + + @mcp.resource("resource://{org}/{repo}/data") + def get_data_mismatched(org: str, repo_2: str) -> str: + return f"Data for {org}" + + """Test that a resource with no parameters works as a regular resource""" + mcp = FastMCP() + + @mcp.resource("resource://static") + def get_static_data() -> str: + return "Static data" + + async with client_session(mcp._mcp_server) as client: + result = await client.read_resource(AnyUrl("resource://static")) + assert isinstance(result.contents[0], TextResourceContents) + assert result.contents[0].text == "Static data" + + @pytest.mark.anyio + async def test_template_to_resource_conversion(self): + """Test that templates are properly converted to resources when accessed""" + mcp = FastMCP() + + @mcp.resource("resource://{name}/data") + def get_data(name: str) -> str: + return f"Data for {name}" + + # Should be registered as a template + assert len(mcp._resource_manager._templates) == 1 + assert len(await mcp.list_resources()) == 0 + + # When accessed, should create a concrete resource + resource = await mcp._resource_manager.get_resource("resource://test/data") + assert isinstance(resource, FunctionResource) + result = await resource.read() + assert result == "Data for test" + + +class TestContextInjection: + """Test context injection in tools.""" + + @pytest.mark.anyio + async def test_context_detection(self): + """Test that context parameters are properly detected.""" + mcp = FastMCP() + + def tool_with_context(x: int, ctx: Context) -> str: + return f"Request {ctx.request_id}: {x}" + + tool = mcp._tool_manager.add_tool(tool_with_context) + assert tool.context_kwarg == "ctx" + + @pytest.mark.anyio + async def test_context_injection(self): + """Test that context is properly injected into tool calls.""" + mcp = FastMCP() + + def tool_with_context(x: int, ctx: Context) -> str: + assert ctx.request_id is not None + return f"Request {ctx.request_id}: {x}" + + mcp.add_tool(tool_with_context) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("tool_with_context", {"x": 42}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert "Request" in content.text + assert "42" in content.text + + @pytest.mark.anyio + async def test_async_context(self): + """Test that context works in async functions.""" + mcp = FastMCP() + + async def async_tool(x: int, ctx: Context) -> str: + assert ctx.request_id is not None + return f"Async request {ctx.request_id}: {x}" + + mcp.add_tool(async_tool) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("async_tool", {"x": 42}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert "Async request" in content.text + assert "42" in content.text + + @pytest.mark.anyio + async def test_context_logging(self): + from unittest.mock import patch + + import mcp.server.session + + """Test that context logging methods work.""" + mcp = FastMCP() + + async def logging_tool(msg: str, ctx: Context) -> str: + await ctx.debug("Debug message") + await ctx.info("Info message") + await ctx.warning("Warning message") + await ctx.error("Error message") + return f"Logged messages for {msg}" + + mcp.add_tool(logging_tool) + + with patch("mcp.server.session.ServerSession.send_log_message") as mock_log: + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("logging_tool", {"msg": "test"}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert "Logged messages for test" in content.text + + assert mock_log.call_count == 4 + mock_log.assert_any_call( + level="debug", data="Debug message", logger=None + ) + mock_log.assert_any_call(level="info", data="Info message", logger=None) + mock_log.assert_any_call( + level="warning", data="Warning message", logger=None + ) + mock_log.assert_any_call( + level="error", data="Error message", logger=None + ) + + @pytest.mark.anyio + async def test_optional_context(self): + """Test that context is optional.""" + mcp = FastMCP() + + def no_context(x: int) -> int: + return x * 2 + + mcp.add_tool(no_context) + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("no_context", {"x": 21}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert content.text == "42" + + @pytest.mark.anyio + async def test_context_resource_access(self): + """Test that context can access resources.""" + mcp = FastMCP() + + @mcp.resource("test://data") + def test_resource() -> str: + return "resource data" + + @mcp.tool() + async def tool_with_resource(ctx: Context) -> str: + r_iter = await ctx.read_resource("test://data") + r_list = list(r_iter) + assert len(r_list) == 1 + r = r_list[0] + return f"Read resource: {r.content} with mime type {r.mime_type}" + + async with client_session(mcp._mcp_server) as client: + result = await client.call_tool("tool_with_resource", {}) + assert len(result.content) == 1 + content = result.content[0] + assert isinstance(content, TextContent) + assert "Read resource: resource data" in content.text + + +class TestServerPrompts: + """Test prompt functionality in FastMCP server.""" + + @pytest.mark.anyio + async def test_prompt_decorator(self): + """Test that the prompt decorator registers prompts correctly.""" + mcp = FastMCP() + + @mcp.prompt() + def fn() -> str: + return "Hello, world!" + + prompts = mcp._prompt_manager.list_prompts() + assert len(prompts) == 1 + assert prompts[0].name == "fn" + # Don't compare functions directly since validate_call wraps them + content = await prompts[0].render() + assert isinstance(content[0].content, TextContent) + assert content[0].content.text == "Hello, world!" + + @pytest.mark.anyio + async def test_prompt_decorator_with_name(self): + """Test prompt decorator with custom name.""" + mcp = FastMCP() + + @mcp.prompt(name="custom_name") + def fn() -> str: + return "Hello, world!" + + prompts = mcp._prompt_manager.list_prompts() + assert len(prompts) == 1 + assert prompts[0].name == "custom_name" + content = await prompts[0].render() + assert isinstance(content[0].content, TextContent) + assert content[0].content.text == "Hello, world!" + + @pytest.mark.anyio + async def test_prompt_decorator_with_description(self): + """Test prompt decorator with custom description.""" + mcp = FastMCP() + + @mcp.prompt(description="A custom description") + def fn() -> str: + return "Hello, world!" + + prompts = mcp._prompt_manager.list_prompts() + assert len(prompts) == 1 + assert prompts[0].description == "A custom description" + content = await prompts[0].render() + assert isinstance(content[0].content, TextContent) + assert content[0].content.text == "Hello, world!" + + def test_prompt_decorator_error(self): + """Test error when decorator is used incorrectly.""" + mcp = FastMCP() + with pytest.raises(TypeError, match="decorator was used incorrectly"): + + @mcp.prompt # type: ignore + def fn() -> str: + return "Hello, world!" + + @pytest.mark.anyio + async def test_list_prompts(self): + """Test listing prompts through MCP protocol.""" + mcp = FastMCP() + + @mcp.prompt() + def fn(name: str, optional: str = "default") -> str: + return f"Hello, {name}!" + + async with client_session(mcp._mcp_server) as client: + result = await client.list_prompts() + assert result.prompts is not None + assert len(result.prompts) == 1 + prompt = result.prompts[0] + assert prompt.name == "fn" + assert prompt.arguments is not None + assert len(prompt.arguments) == 2 + assert prompt.arguments[0].name == "name" + assert prompt.arguments[0].required is True + assert prompt.arguments[1].name == "optional" + assert prompt.arguments[1].required is False + + @pytest.mark.anyio + async def test_get_prompt(self): + """Test getting a prompt through MCP protocol.""" + mcp = FastMCP() + + @mcp.prompt() + def fn(name: str) -> str: + return f"Hello, {name}!" + + async with client_session(mcp._mcp_server) as client: + result = await client.get_prompt("fn", {"name": "World"}) + assert len(result.messages) == 1 + message = result.messages[0] + assert message.role == "user" + content = message.content + assert isinstance(content, TextContent) + assert content.text == "Hello, World!" + + @pytest.mark.anyio + async def test_get_prompt_with_resource(self): + """Test getting a prompt that returns resource content.""" + mcp = FastMCP() + + @mcp.prompt() + def fn() -> Message: + return UserMessage( + content=EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri=AnyUrl("file://file.txt"), + text="File contents", + mimeType="text/plain", + ), + ) + ) + + async with client_session(mcp._mcp_server) as client: + result = await client.get_prompt("fn") + assert len(result.messages) == 1 + message = result.messages[0] + assert message.role == "user" + content = message.content + assert isinstance(content, EmbeddedResource) + resource = content.resource + assert isinstance(resource, TextResourceContents) + assert resource.text == "File contents" + assert resource.mimeType == "text/plain" + + @pytest.mark.anyio + async def test_get_unknown_prompt(self): + """Test error when getting unknown prompt.""" + mcp = FastMCP() + async with client_session(mcp._mcp_server) as client: + with pytest.raises(McpError, match="Unknown prompt"): + await client.get_prompt("unknown") + + @pytest.mark.anyio + async def test_get_prompt_missing_args(self): + """Test error when required arguments are missing.""" + mcp = FastMCP() + + @mcp.prompt() + def prompt_fn(name: str) -> str: + return f"Hello, {name}!" + + async with client_session(mcp._mcp_server) as client: + with pytest.raises(McpError, match="Missing required arguments"): + await client.get_prompt("prompt_fn") diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index 3accb8178..4115b00c9 100644 --- a/tests/tools/test_tool_manager.py +++ b/tests/tools/test_tool_manager.py @@ -1,101 +1,431 @@ -from fastmcp.tools.tool_manager import ToolManager +import json +import logging + +import pytest +from pydantic import BaseModel + +from fastmcp.exceptions import ToolError +from fastmcp.tools import ToolManager -def test_import_tools(): - """Test importing tools from one manager to another with a prefix.""" - # Setup source manager with tools - source_manager = ToolManager() +class TestAddTools: + def test_basic_function(self): + """Test registering and running a basic function.""" - # Create some test tools - def tool1_fn(): - return "Tool 1 result" + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b - def tool2_fn(): - return "Tool 2 result" + manager = ToolManager() + manager.add_tool(add) - # Add tools to source manager - source_manager.add_tool(tool1_fn, name="get_data", description="Get some data") - source_manager.add_tool( - tool2_fn, name="process_data", description="Process the data" - ) + tool = manager.get_tool("add") + assert tool is not None + assert tool.name == "add" + assert tool.description == "Add two numbers." + assert tool.is_async is False + assert tool.parameters["properties"]["a"]["type"] == "integer" + assert tool.parameters["properties"]["b"]["type"] == "integer" - # Create target manager - target_manager = ToolManager() + @pytest.mark.anyio + async def test_async_function(self): + """Test registering and running an async function.""" - # Import tools from source to target - prefix = "source/" - target_manager.import_tools(source_manager, prefix) + async def fetch_data(url: str) -> str: + """Fetch data from URL.""" + return f"Data from {url}" - # Verify tools were imported with prefixes - assert "source/get_data" in target_manager._tools - assert "source/process_data" in target_manager._tools + manager = ToolManager() + manager.add_tool(fetch_data) - # Verify the original tools still exist in source manager - assert "get_data" in source_manager._tools - assert "process_data" in source_manager._tools + tool = manager.get_tool("fetch_data") + assert tool is not None + assert tool.name == "fetch_data" + assert tool.description == "Fetch data from URL." + assert tool.is_async is True + assert tool.parameters["properties"]["url"]["type"] == "string" - # Verify the imported tools have the correct descriptions - assert target_manager._tools["source/get_data"].description == "Get some data" - assert ( - target_manager._tools["source/process_data"].description == "Process the data" - ) + def test_pydantic_model_function(self): + """Test registering a function that takes a Pydantic model.""" - # Verify the tool functions were properly copied - # We can't directly compare functions, so we'll check their __name__ attribute - assert target_manager._tools["source/get_data"].fn.__name__ == tool1_fn.__name__ - assert target_manager._tools["source/process_data"].fn.__name__ == tool2_fn.__name__ + class UserInput(BaseModel): + name: str + age: int + + def create_user(user: UserInput, flag: bool) -> dict: + """Create a new user.""" + return {"id": 1, **user.model_dump()} + + manager = ToolManager() + manager.add_tool(create_user) + + tool = manager.get_tool("create_user") + assert tool is not None + assert tool.name == "create_user" + assert tool.description == "Create a new user." + assert tool.is_async is False + assert "name" in tool.parameters["$defs"]["UserInput"]["properties"] + assert "age" in tool.parameters["$defs"]["UserInput"]["properties"] + assert "flag" in tool.parameters["properties"] + + def test_add_invalid_tool(self): + manager = ToolManager() + with pytest.raises(AttributeError): + manager.add_tool(1) # type: ignore + + def test_add_lambda(self): + manager = ToolManager() + tool = manager.add_tool(lambda x: x, name="my_tool") + assert tool.name == "my_tool" + + def test_add_lambda_with_no_name(self): + manager = ToolManager() + with pytest.raises( + ValueError, match="You must provide a name for lambda functions" + ): + manager.add_tool(lambda x: x) + + def test_warn_on_duplicate_tools(self, caplog): + """Test warning on duplicate tools.""" + + def f(x: int) -> int: + return x + + manager = ToolManager() + manager.add_tool(f) + with caplog.at_level(logging.WARNING): + manager.add_tool(f) + assert "Tool already exists: f" in caplog.text + + def test_disable_warn_on_duplicate_tools(self, caplog): + """Test disabling warning on duplicate tools.""" + + def f(x: int) -> int: + return x + + manager = ToolManager() + manager.add_tool(f) + manager.warn_on_duplicate_tools = False + with caplog.at_level(logging.WARNING): + manager.add_tool(f) + assert "Tool already exists: f" not in caplog.text -def test_tool_duplicate_behavior(): - """Test the behavior when importing tools with duplicate names.""" - # Setup source and target managers - source_manager = ToolManager() - target_manager = ToolManager() +class TestCallTools: + @pytest.mark.anyio + async def test_call_tool(self): + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b - # Add the same tool name to both managers - def source_fn(): - return "Source result" + manager = ToolManager() + manager.add_tool(add) + result = await manager.call_tool("add", {"a": 1, "b": 2}) + assert result == 3 - def target_fn(): - return "Target result" + @pytest.mark.anyio + async def test_call_async_tool(self): + async def double(n: int) -> int: + """Double a number.""" + return n * 2 - source_manager.add_tool(source_fn, name="common_tool") - target_manager.add_tool( - target_fn, name="source/common_tool" - ) # Pre-create with the prefixed name + manager = ToolManager() + manager.add_tool(double) + result = await manager.call_tool("double", {"n": 5}) + assert result == 10 - # Import tools from source to target - target_manager.import_tools(source_manager, "source/") + @pytest.mark.anyio + async def test_call_tool_with_default_args(self): + def add(a: int, b: int = 1) -> int: + """Add two numbers.""" + return a + b - # The original tool in the target manager is replaced by the imported one - assert target_manager._tools["source/common_tool"].fn.__name__ == source_fn.__name__ + manager = ToolManager() + manager.add_tool(add) + result = await manager.call_tool("add", {"a": 1}) + assert result == 2 + + @pytest.mark.anyio + async def test_call_tool_with_missing_args(self): + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + manager = ToolManager() + manager.add_tool(add) + with pytest.raises(ToolError): + await manager.call_tool("add", {"a": 1}) + + @pytest.mark.anyio + async def test_call_unknown_tool(self): + manager = ToolManager() + with pytest.raises(ToolError): + await manager.call_tool("unknown", {"a": 1}) + + @pytest.mark.anyio + async def test_call_tool_with_list_int_input(self): + def sum_vals(vals: list[int]) -> int: + return sum(vals) + + manager = ToolManager() + manager.add_tool(sum_vals) + # Try both with plain list and with JSON list + result = await manager.call_tool("sum_vals", {"vals": "[1, 2, 3]"}) + assert result == 6 + result = await manager.call_tool("sum_vals", {"vals": [1, 2, 3]}) + assert result == 6 + + @pytest.mark.anyio + async def test_call_tool_with_list_str_or_str_input(self): + def concat_strs(vals: list[str] | str) -> str: + return vals if isinstance(vals, str) else "".join(vals) + + manager = ToolManager() + manager.add_tool(concat_strs) + # Try both with plain python object and with JSON list + result = await manager.call_tool("concat_strs", {"vals": ["a", "b", "c"]}) + assert result == "abc" + result = await manager.call_tool("concat_strs", {"vals": '["a", "b", "c"]'}) + assert result == "abc" + result = await manager.call_tool("concat_strs", {"vals": "a"}) + assert result == "a" + result = await manager.call_tool("concat_strs", {"vals": '"a"'}) + assert result == '"a"' + + @pytest.mark.anyio + async def test_call_tool_with_complex_model(self): + from fastmcp import Context + + class MyShrimpTank(BaseModel): + class Shrimp(BaseModel): + name: str + + shrimp: list[Shrimp] + x: None + + def name_shrimp(tank: MyShrimpTank, ctx: Context) -> list[str]: + return [x.name for x in tank.shrimp] + + manager = ToolManager() + manager.add_tool(name_shrimp) + result = await manager.call_tool( + "name_shrimp", + {"tank": {"x": None, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}}, + ) + assert result == ["rex", "gertrude"] + result = await manager.call_tool( + "name_shrimp", + {"tank": '{"x": null, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}'}, + ) + assert result == ["rex", "gertrude"] -def test_import_tools_with_multiple_prefixes(): - """Test importing tools from multiple managers with different prefixes.""" - # Setup source managers - weather_manager = ToolManager() - news_manager = ToolManager() +class TestToolSchema: + @pytest.mark.anyio + async def test_context_arg_excluded_from_schema(self): + from fastmcp import Context - # Add tools to source managers - def forecast_fn(): - return "Weather forecast" + def something(a: int, ctx: Context) -> int: + return a - def headlines_fn(): - return "News headlines" + manager = ToolManager() + tool = manager.add_tool(something) + assert "ctx" not in json.dumps(tool.parameters) + assert "Context" not in json.dumps(tool.parameters) + assert "ctx" not in tool.fn_metadata.arg_model.model_fields - weather_manager.add_tool(forecast_fn, name="forecast") - news_manager.add_tool(headlines_fn, name="headlines") - # Create target manager and import from both sources - main_manager = ToolManager() - main_manager.import_tools(weather_manager, "weather/") - main_manager.import_tools(news_manager, "news/") +class TestContextHandling: + """Test context handling in the tool manager.""" - # Verify tools were imported with correct prefixes - assert "weather/forecast" in main_manager._tools - assert "news/headlines" in main_manager._tools + def test_context_parameter_detection(self): + """Test that context parameters are properly detected in + Tool.from_function().""" + from fastmcp import Context - # Verify the tools are accessible and functioning - assert main_manager._tools["weather/forecast"].fn.__name__ == forecast_fn.__name__ - assert main_manager._tools["news/headlines"].fn.__name__ == headlines_fn.__name__ + def tool_with_context(x: int, ctx: Context) -> str: + return str(x) + + manager = ToolManager() + tool = manager.add_tool(tool_with_context) + assert tool.context_kwarg == "ctx" + + def tool_without_context(x: int) -> str: + return str(x) + + tool = manager.add_tool(tool_without_context) + assert tool.context_kwarg is None + + @pytest.mark.anyio + async def test_context_injection(self): + """Test that context is properly injected during tool execution.""" + from fastmcp import Context, FastMCP + + def tool_with_context(x: int, ctx: Context) -> str: + assert isinstance(ctx, Context) + return str(x) + + manager = ToolManager() + manager.add_tool(tool_with_context) + + mcp = FastMCP() + ctx = mcp.get_context() + result = await manager.call_tool("tool_with_context", {"x": 42}, context=ctx) + assert result == "42" + + @pytest.mark.anyio + async def test_context_injection_async(self): + """Test that context is properly injected in async tools.""" + from fastmcp import Context, FastMCP + + async def async_tool(x: int, ctx: Context) -> str: + assert isinstance(ctx, Context) + return str(x) + + manager = ToolManager() + manager.add_tool(async_tool) + + mcp = FastMCP() + ctx = mcp.get_context() + result = await manager.call_tool("async_tool", {"x": 42}, context=ctx) + assert result == "42" + + @pytest.mark.anyio + async def test_context_optional(self): + """Test that context is optional when calling tools.""" + from fastmcp import Context + + def tool_with_context(x: int, ctx: Context | None = None) -> str: + return str(x) + + manager = ToolManager() + manager.add_tool(tool_with_context) + # Should not raise an error when context is not provided + result = await manager.call_tool("tool_with_context", {"x": 42}) + assert result == "42" + + @pytest.mark.anyio + async def test_context_error_handling(self): + """Test error handling when context injection fails.""" + from fastmcp import Context, FastMCP + + def tool_with_context(x: int, ctx: Context) -> str: + raise ValueError("Test error") + + manager = ToolManager() + manager.add_tool(tool_with_context) + + mcp = FastMCP() + ctx = mcp.get_context() + with pytest.raises(ToolError, match="Error executing tool tool_with_context"): + await manager.call_tool("tool_with_context", {"x": 42}, context=ctx) + + +class TestImportTools: + def test_import_tools(self): + """Test importing tools from one manager to another with a prefix.""" + # Setup source manager with tools + source_manager = ToolManager() + + # Create some test tools + def tool1_fn(): + return "Tool 1 result" + + def tool2_fn(): + return "Tool 2 result" + + # Add tools to source manager + source_manager.add_tool(tool1_fn, name="get_data", description="Get some data") + source_manager.add_tool( + tool2_fn, name="process_data", description="Process the data" + ) + + # Create target manager + target_manager = ToolManager() + + # Import tools from source to target + prefix = "source/" + target_manager.import_tools(source_manager, prefix) + + # Verify tools were imported with prefixes + assert "source/get_data" in target_manager._tools + assert "source/process_data" in target_manager._tools + + # Verify the original tools still exist in source manager + assert "get_data" in source_manager._tools + assert "process_data" in source_manager._tools + + # Verify the imported tools have the correct descriptions + assert target_manager._tools["source/get_data"].description == "Get some data" + assert ( + target_manager._tools["source/process_data"].description + == "Process the data" + ) + + # Verify the tool functions were properly copied + # We can't directly compare functions, so we'll check their __name__ attribute + assert target_manager._tools["source/get_data"].fn.__name__ == tool1_fn.__name__ + assert ( + target_manager._tools["source/process_data"].fn.__name__ + == tool2_fn.__name__ + ) + + def test_tool_duplicate_behavior(self): + """Test the behavior when importing tools with duplicate names.""" + # Setup source and target managers + source_manager = ToolManager() + target_manager = ToolManager() + + # Add the same tool name to both managers + def source_fn(): + return "Source result" + + def target_fn(): + return "Target result" + + source_manager.add_tool(source_fn, name="common_tool") + target_manager.add_tool( + target_fn, name="source/common_tool" + ) # Pre-create with the prefixed name + + # Import tools from source to target + target_manager.import_tools(source_manager, "source/") + + # The original tool in the target manager is replaced by the imported one + assert ( + target_manager._tools["source/common_tool"].fn.__name__ + == source_fn.__name__ + ) + + def test_import_tools_with_multiple_prefixes(self): + """Test importing tools from multiple managers with different prefixes.""" + # Setup source managers + weather_manager = ToolManager() + news_manager = ToolManager() + + # Add tools to source managers + def forecast_fn(): + return "Weather forecast" + + def headlines_fn(): + return "News headlines" + + weather_manager.add_tool(forecast_fn, name="forecast") + news_manager.add_tool(headlines_fn, name="headlines") + + # Create target manager and import from both sources + main_manager = ToolManager() + main_manager.import_tools(weather_manager, "weather/") + main_manager.import_tools(news_manager, "news/") + + # Verify tools were imported with correct prefixes + assert "weather/forecast" in main_manager._tools + assert "news/headlines" in main_manager._tools + + # Verify the tools are accessible and functioning + assert ( + main_manager._tools["weather/forecast"].fn.__name__ == forecast_fn.__name__ + ) + assert ( + main_manager._tools["news/headlines"].fn.__name__ == headlines_fn.__name__ + ) diff --git a/tests/tools/tool_manager.py b/tests/tools/tool_manager.py deleted file mode 100644 index 78da44b37..000000000 --- a/tests/tools/tool_manager.py +++ /dev/null @@ -1,104 +0,0 @@ -from fastmcp.tools.tool_manager import ToolManager - - -def test_import_tools(): - """Test importing tools from one manager to another with a prefix.""" - # Setup source manager with tools - source_manager = ToolManager() - - # Create some test tools - def tool1_fn(): - return "Tool 1 result" - - def tool2_fn(): - return "Tool 2 result" - - # Add tools to source manager - source_manager.add_tool(tool1_fn, name="get_data", description="Get some data") - source_manager.add_tool( - tool2_fn, name="process_data", description="Process the data" - ) - - # Create target manager - target_manager = ToolManager() - - # Import tools from source to target - prefix = "source" - target_manager.import_tools(source_manager, prefix) - - # Verify tools were imported with prefixes - assert "source:get_data" in target_manager._tools - assert "source:process_data" in target_manager._tools - - # Verify the original tools still exist in source manager - assert "get_data" in source_manager._tools - assert "process_data" in source_manager._tools - - # Verify the imported tools have the correct descriptions - assert target_manager._tools["source:get_data"].description == "Get some data" - assert ( - target_manager._tools["source:process_data"].description == "Process the data" - ) - - # Verify the tool functions were properly copied - # We can't directly compare functions, so we'll check their __name__ attribute - assert target_manager._tools["source:get_data"].fn.__name__ == tool1_fn.__name__ - assert target_manager._tools["source:process_data"].fn.__name__ == tool2_fn.__name__ - - -def test_import_tools_duplicate_warning(caplog): - """Test that warning is logged when importing a tool with a name that already exists.""" - # Setup source and target managers - source_manager = ToolManager() - target_manager = ToolManager(warn_on_duplicate_tools=True) - - # Add the same tool name to both managers - def source_fn(): - return "Source result" - - def target_fn(): - return "Target result" - - source_manager.add_tool(source_fn, name="common_tool") - target_manager.add_tool( - target_fn, name="source:common_tool" - ) # Pre-create with the prefixed name - - # Import tools from source to target - target_manager.import_tools(source_manager, "source") - - # Verify a warning was logged - assert any("already exists" in record.message for record in caplog.records) - - # The original tool in the target manager should be preserved - assert target_manager._tools["source:common_tool"].fn.__name__ == target_fn.__name__ - - -def test_import_tools_with_multiple_prefixes(): - """Test importing tools from multiple managers with different prefixes.""" - # Setup source managers - weather_manager = ToolManager() - news_manager = ToolManager() - - # Add tools to source managers - def forecast_fn(): - return "Weather forecast" - - def headlines_fn(): - return "News headlines" - - weather_manager.add_tool(forecast_fn, name="forecast") - news_manager.add_tool(headlines_fn, name="headlines") - - # Create target manager and import from both sources - main_manager = ToolManager() - main_manager.import_tools(weather_manager, "weather") - main_manager.import_tools(news_manager, "news") - - # Verify tools were imported with correct prefixes - assert "weather:forecast" in main_manager._tools - assert "news:headlines" in main_manager._tools - - # Verify the tools are accessible and functioning - assert main_manager._tools["weather:forecast"].fn.__name__ == forecast_fn.__name__ - assert main_manager._tools["news:headlines"].fn.__name__ == headlines_fn.__name__ diff --git a/tests/utilities/test_func_metadata.py b/tests/utilities/test_func_metadata.py new file mode 100644 index 000000000..ee8037c7f --- /dev/null +++ b/tests/utilities/test_func_metadata.py @@ -0,0 +1,416 @@ +from typing import Annotated + +import annotated_types +import pytest +from pydantic import BaseModel, Field + +from fastmcp.utilities.func_metadata import func_metadata + + +class SomeInputModelA(BaseModel): + pass + + +class SomeInputModelB(BaseModel): + class InnerModel(BaseModel): + x: int + + how_many_shrimp: Annotated[int, Field(description="How many shrimp in the tank???")] + ok: InnerModel + y: None + + +def complex_arguments_fn( + an_int: int, + must_be_none: None, + must_be_none_dumb_annotation: Annotated[None, "blah"], + list_of_ints: list[int], + # list[str] | str is an interesting case because if it comes in as JSON like + # "[\"a\", \"b\"]" then it will be naively parsed as a string. + list_str_or_str: list[str] | str, + an_int_annotated_with_field: Annotated[ + int, Field(description="An int with a field") + ], + an_int_annotated_with_field_and_others: Annotated[ + int, + str, # Should be ignored, really + Field(description="An int with a field"), + annotated_types.Gt(1), + ], + an_int_annotated_with_junk: Annotated[ + int, + "123", + 456, + ], + field_with_default_via_field_annotation_before_nondefault_arg: Annotated[ + int, Field(1) + ], + unannotated, + my_model_a: SomeInputModelA, + my_model_a_forward_ref: "SomeInputModelA", + my_model_b: SomeInputModelB, + an_int_annotated_with_field_default: Annotated[ + int, + Field(1, description="An int with a field"), + ], + unannotated_with_default=5, + my_model_a_with_default: SomeInputModelA = SomeInputModelA(), # noqa: B008 + an_int_with_default: int = 1, + must_be_none_with_default: None = None, + an_int_with_equals_field: int = Field(1, ge=0), + int_annotated_with_default: Annotated[int, Field(description="hey")] = 5, +) -> str: + _ = ( + an_int, + must_be_none, + must_be_none_dumb_annotation, + list_of_ints, + list_str_or_str, + an_int_annotated_with_field, + an_int_annotated_with_field_and_others, + an_int_annotated_with_junk, + field_with_default_via_field_annotation_before_nondefault_arg, + unannotated, + an_int_annotated_with_field_default, + unannotated_with_default, + my_model_a, + my_model_a_forward_ref, + my_model_b, + my_model_a_with_default, + an_int_with_default, + must_be_none_with_default, + an_int_with_equals_field, + int_annotated_with_default, + ) + return "ok!" + + +@pytest.mark.anyio +async def test_complex_function_runtime_arg_validation_non_json(): + """Test that basic non-JSON arguments are validated correctly""" + meta = func_metadata(complex_arguments_fn) + + # Test with minimum required arguments + result = await meta.call_fn_with_arg_validation( + complex_arguments_fn, + fn_is_async=False, + arguments_to_validate={ + "an_int": 1, + "must_be_none": None, + "must_be_none_dumb_annotation": None, + "list_of_ints": [1, 2, 3], + "list_str_or_str": "hello", + "an_int_annotated_with_field": 42, + "an_int_annotated_with_field_and_others": 5, + "an_int_annotated_with_junk": 100, + "unannotated": "test", + "my_model_a": {}, + "my_model_a_forward_ref": {}, + "my_model_b": {"how_many_shrimp": 5, "ok": {"x": 1}, "y": None}, + }, + arguments_to_pass_directly=None, + ) + assert result == "ok!" + + # Test with invalid types + with pytest.raises(ValueError): + await meta.call_fn_with_arg_validation( + complex_arguments_fn, + fn_is_async=False, + arguments_to_validate={"an_int": "not an int"}, + arguments_to_pass_directly=None, + ) + + +@pytest.mark.anyio +async def test_complex_function_runtime_arg_validation_with_json(): + """Test that JSON string arguments are parsed and validated correctly""" + meta = func_metadata(complex_arguments_fn) + + result = await meta.call_fn_with_arg_validation( + complex_arguments_fn, + fn_is_async=False, + arguments_to_validate={ + "an_int": 1, + "must_be_none": None, + "must_be_none_dumb_annotation": None, + "list_of_ints": "[1, 2, 3]", # JSON string + "list_str_or_str": '["a", "b", "c"]', # JSON string + "an_int_annotated_with_field": 42, + "an_int_annotated_with_field_and_others": "5", # JSON string + "an_int_annotated_with_junk": 100, + "unannotated": "test", + "my_model_a": "{}", # JSON string + "my_model_a_forward_ref": "{}", # JSON string + "my_model_b": '{"how_many_shrimp": 5, "ok": {"x": 1}, "y": null}', + }, + arguments_to_pass_directly=None, + ) + assert result == "ok!" + + +def test_str_vs_list_str(): + """Test handling of string vs list[str] type annotations. + + This is tricky as '"hello"' can be parsed as a JSON string or a Python string. + We want to make sure it's kept as a python string. + """ + + def func_with_str_types(str_or_list: str | list[str]): + return str_or_list + + meta = func_metadata(func_with_str_types) + + # Test string input for union type + result = meta.pre_parse_json({"str_or_list": "hello"}) + assert result["str_or_list"] == "hello" + + # Test string input that contains valid JSON for union type + # We want to see here that the JSON-vali string is NOT parsed as JSON, but rather + # kept as a raw string + result = meta.pre_parse_json({"str_or_list": '"hello"'}) + assert result["str_or_list"] == '"hello"' + + # Test list input for union type + result = meta.pre_parse_json({"str_or_list": '["hello", "world"]'}) + assert result["str_or_list"] == ["hello", "world"] + + +def test_skip_names(): + """Test that skipped parameters are not included in the model""" + + def func_with_many_params( + keep_this: int, skip_this: str, also_keep: float, also_skip: bool + ): + return keep_this, skip_this, also_keep, also_skip + + # Skip some parameters + meta = func_metadata(func_with_many_params, skip_names=["skip_this", "also_skip"]) + + # Check model fields + assert "keep_this" in meta.arg_model.model_fields + assert "also_keep" in meta.arg_model.model_fields + assert "skip_this" not in meta.arg_model.model_fields + assert "also_skip" not in meta.arg_model.model_fields + + # Validate that we can call with only non-skipped parameters + model: BaseModel = meta.arg_model.model_validate({"keep_this": 1, "also_keep": 2.5}) # type: ignore + assert model.keep_this == 1 # type: ignore + assert model.also_keep == 2.5 # type: ignore + + +@pytest.mark.anyio +async def test_lambda_function(): + """Test lambda function schema and validation""" + fn = lambda x, y=5: x # noqa: E731 + meta = func_metadata(lambda x, y=5: x) + + # Test schema + assert meta.arg_model.model_json_schema() == { + "properties": { + "x": {"title": "x", "type": "string"}, + "y": {"default": 5, "title": "y", "type": "string"}, + }, + "required": ["x"], + "title": "Arguments", + "type": "object", + } + + async def check_call(args): + return await meta.call_fn_with_arg_validation( + fn, + fn_is_async=False, + arguments_to_validate=args, + arguments_to_pass_directly=None, + ) + + # Basic calls + assert await check_call({"x": "hello"}) == "hello" + assert await check_call({"x": "hello", "y": "world"}) == "hello" + assert await check_call({"x": '"hello"'}) == '"hello"' + + # Missing required arg + with pytest.raises(ValueError): + await check_call({"y": "world"}) + + +def test_complex_function_json_schema(): + """Test JSON schema generation for complex function arguments. + + Note: Different versions of pydantic output slightly different + JSON Schema formats for model fields with defaults. The format changed in 2.9.0: + + 1. Before 2.9.0: + { + "allOf": [{"$ref": "#/$defs/Model"}], + "default": {} + } + + 2. Since 2.9.0: + { + "$ref": "#/$defs/Model", + "default": {} + } + + Both formats are valid and functionally equivalent. This test accepts either format + to ensure compatibility across our supported pydantic versions. + + This change in format does not affect runtime behavior since: + 1. Both schemas validate the same way + 2. The actual model classes and validation logic are unchanged + 3. func_metadata uses model_validate/model_dump, not the schema directly + """ + meta = func_metadata(complex_arguments_fn) + actual_schema = meta.arg_model.model_json_schema() + + # Create a copy of the actual schema to normalize + normalized_schema = actual_schema.copy() + + # Normalize the my_model_a_with_default field to handle both pydantic formats + if "allOf" in actual_schema["properties"]["my_model_a_with_default"]: + normalized_schema["properties"]["my_model_a_with_default"] = { + "$ref": "#/$defs/SomeInputModelA", + "default": {}, + } + + assert normalized_schema == { + "$defs": { + "InnerModel": { + "properties": {"x": {"title": "X", "type": "integer"}}, + "required": ["x"], + "title": "InnerModel", + "type": "object", + }, + "SomeInputModelA": { + "properties": {}, + "title": "SomeInputModelA", + "type": "object", + }, + "SomeInputModelB": { + "properties": { + "how_many_shrimp": { + "description": "How many shrimp in the tank???", + "title": "How Many Shrimp", + "type": "integer", + }, + "ok": {"$ref": "#/$defs/InnerModel"}, + "y": {"title": "Y", "type": "null"}, + }, + "required": ["how_many_shrimp", "ok", "y"], + "title": "SomeInputModelB", + "type": "object", + }, + }, + "properties": { + "an_int": {"title": "An Int", "type": "integer"}, + "must_be_none": {"title": "Must Be None", "type": "null"}, + "must_be_none_dumb_annotation": { + "title": "Must Be None Dumb Annotation", + "type": "null", + }, + "list_of_ints": { + "items": {"type": "integer"}, + "title": "List Of Ints", + "type": "array", + }, + "list_str_or_str": { + "anyOf": [ + {"items": {"type": "string"}, "type": "array"}, + {"type": "string"}, + ], + "title": "List Str Or Str", + }, + "an_int_annotated_with_field": { + "description": "An int with a field", + "title": "An Int Annotated With Field", + "type": "integer", + }, + "an_int_annotated_with_field_and_others": { + "description": "An int with a field", + "exclusiveMinimum": 1, + "title": "An Int Annotated With Field And Others", + "type": "integer", + }, + "an_int_annotated_with_junk": { + "title": "An Int Annotated With Junk", + "type": "integer", + }, + "field_with_default_via_field_annotation_before_nondefault_arg": { + "default": 1, + "title": "Field With Default Via Field Annotation Before Nondefault Arg", + "type": "integer", + }, + "unannotated": {"title": "unannotated", "type": "string"}, + "my_model_a": {"$ref": "#/$defs/SomeInputModelA"}, + "my_model_a_forward_ref": {"$ref": "#/$defs/SomeInputModelA"}, + "my_model_b": {"$ref": "#/$defs/SomeInputModelB"}, + "an_int_annotated_with_field_default": { + "default": 1, + "description": "An int with a field", + "title": "An Int Annotated With Field Default", + "type": "integer", + }, + "unannotated_with_default": { + "default": 5, + "title": "unannotated_with_default", + "type": "string", + }, + "my_model_a_with_default": { + "$ref": "#/$defs/SomeInputModelA", + "default": {}, + }, + "an_int_with_default": { + "default": 1, + "title": "An Int With Default", + "type": "integer", + }, + "must_be_none_with_default": { + "default": None, + "title": "Must Be None With Default", + "type": "null", + }, + "an_int_with_equals_field": { + "default": 1, + "minimum": 0, + "title": "An Int With Equals Field", + "type": "integer", + }, + "int_annotated_with_default": { + "default": 5, + "description": "hey", + "title": "Int Annotated With Default", + "type": "integer", + }, + }, + "required": [ + "an_int", + "must_be_none", + "must_be_none_dumb_annotation", + "list_of_ints", + "list_str_or_str", + "an_int_annotated_with_field", + "an_int_annotated_with_field_and_others", + "an_int_annotated_with_junk", + "unannotated", + "my_model_a", + "my_model_a_forward_ref", + "my_model_b", + ], + "title": "complex_arguments_fnArguments", + "type": "object", + } + + +def test_str_vs_int(): + """ + Test that string values are kept as strings even when they contain numbers, + while numbers are parsed correctly. + """ + + def func_with_str_and_int(a: str, b: int): + return a + + meta = func_metadata(func_with_str_and_int) + result = meta.pre_parse_json({"a": "123", "b": 123}) + assert result["a"] == "123" + assert result["b"] == 123