fastmcp 0.4.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

fastmcp/client/transports.py

@@ -0,0 +1,411 @@
+import abc
+import contextlib
+import datetime
+import os
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import (
+    TypedDict,
+)
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.session import (
+    ListRootsFnT,
+    LoggingFnT,
+    MessageHandlerFnT,
+    SamplingFnT,
+)
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from mcp.client.websocket import websocket_client
+from mcp.shared.memory import create_connected_server_and_client_session
+from pydantic import AnyUrl
+from typing_extensions import Unpack
+
+from fastmcp.server import FastMCP as FastMCPServer
+
+
+class SessionKwargs(TypedDict, total=False):
+    """Keyword arguments for the MCP ClientSession constructor."""
+
+    sampling_callback: SamplingFnT | None
+    list_roots_callback: ListRootsFnT | None
+    logging_callback: LoggingFnT | None
+    message_handler: MessageHandlerFnT | None
+    read_timeout_seconds: datetime.timedelta | None
+
+
+class ClientTransport(abc.ABC):
+    """
+    Abstract base class for different MCP client transport mechanisms.
+
+    A Transport is responsible for establishing and managing connections
+    to an MCP server, and providing a ClientSession within an async context.
+    """
+
+    @abc.abstractmethod
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        """
+        Establishes a connection and yields an active, initialized ClientSession.
+
+        The session is guaranteed to be valid only within the scope of the
+        async context manager. Connection setup and teardown are handled
+        within this context.
+
+        Args:
+            **session_kwargs: Keyword arguments to pass to the ClientSession
+                              constructor (e.g., callbacks, timeouts).
+
+        Yields:
+            An initialized mcp.ClientSession instance.
+        """
+        raise NotImplementedError
+        yield None  # type: ignore
+
+    def __repr__(self) -> str:
+        # Basic representation for subclasses
+        return f"<{self.__class__.__name__}>"
+
+
+class WSTransport(ClientTransport):
+    """Transport implementation that connects to an MCP server via WebSockets."""
+
+    def __init__(self, url: str | AnyUrl):
+        if isinstance(url, AnyUrl):
+            url = str(url)
+        if not isinstance(url, str) or not url.startswith("ws"):
+            raise ValueError("Invalid WebSocket URL provided.")
+        self.url = url
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        async with websocket_client(self.url) as transport:
+            read_stream, write_stream = transport
+            async with ClientSession(
+                read_stream, write_stream, **session_kwargs
+            ) as session:
+                await session.initialize()  # Initialize after session creation
+                yield session
+
+    def __repr__(self) -> str:
+        return f"<WebSocket(url='{self.url}')>"
+
+
+class SSETransport(ClientTransport):
+    """Transport implementation that connects to an MCP server via Server-Sent Events."""
+
+    def __init__(self, url: str | AnyUrl, headers: dict[str, str] | None = None):
+        if isinstance(url, AnyUrl):
+            url = str(url)
+        if not isinstance(url, str) or not url.startswith("http"):
+            raise ValueError("Invalid HTTP/S URL provided for SSE.")
+        self.url = url
+        self.headers = headers or {}
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        async with sse_client(self.url, headers=self.headers) as transport:
+            read_stream, write_stream = transport
+            async with ClientSession(
+                read_stream, write_stream, **session_kwargs
+            ) as session:
+                await session.initialize()
+                yield session
+
+    def __repr__(self) -> str:
+        return f"<SSE(url='{self.url}')>"
+
+
+class StdioTransport(ClientTransport):
+    """
+    Base transport for connecting to an MCP server via subprocess with stdio.
+
+    This is a base class that can be subclassed for specific command-based
+    transports like Python, Node, Uvx, etc.
+    """
+
+    def __init__(
+        self,
+        command: str,
+        args: list[str],
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+    ):
+        """
+        Initialize a Stdio transport.
+
+        Args:
+            command: The command to run (e.g., "python", "node", "uvx")
+            args: The arguments to pass to the command
+            env: Environment variables to set for the subprocess
+            cwd: Current working directory for the subprocess
+        """
+        self.command = command
+        self.args = args
+        self.env = env
+        self.cwd = cwd
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        server_params = StdioServerParameters(
+            command=self.command, args=self.args, env=self.env, cwd=self.cwd
+        )
+        async with stdio_client(server_params) as transport:
+            read_stream, write_stream = transport
+            async with ClientSession(
+                read_stream, write_stream, **session_kwargs
+            ) as session:
+                await session.initialize()
+                yield session
+
+    def __repr__(self) -> str:
+        return (
+            f"<{self.__class__.__name__}(command='{self.command}', args={self.args})>"
+        )
+
+
+class PythonStdioTransport(StdioTransport):
+    """Transport for running Python scripts."""
+
+    def __init__(
+        self,
+        script_path: str | Path,
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+        python_cmd: str = "python",
+    ):
+        """
+        Initialize a Python transport.
+
+        Args:
+            script_path: Path to the Python script to run
+            args: Additional arguments to pass to the script
+            env: Environment variables to set for the subprocess
+            cwd: Current working directory for the subprocess
+            python_cmd: Python command to use (default: "python")
+        """
+        script_path = Path(script_path).resolve()
+        if not script_path.is_file():
+            raise FileNotFoundError(f"Script not found: {script_path}")
+        if not str(script_path).endswith(".py"):
+            raise ValueError(f"Not a Python script: {script_path}")
+
+        full_args = [str(script_path)]
+        if args:
+            full_args.extend(args)
+
+        super().__init__(command=python_cmd, args=full_args, env=env, cwd=cwd)
+        self.script_path = script_path
+
+
+class NodeStdioTransport(StdioTransport):
+    """Transport for running Node.js scripts."""
+
+    def __init__(
+        self,
+        script_path: str | Path,
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+        node_cmd: str = "node",
+    ):
+        """
+        Initialize a Node transport.
+
+        Args:
+            script_path: Path to the Node.js script to run
+            args: Additional arguments to pass to the script
+            env: Environment variables to set for the subprocess
+            cwd: Current working directory for the subprocess
+            node_cmd: Node.js command to use (default: "node")
+        """
+        script_path = Path(script_path).resolve()
+        if not script_path.is_file():
+            raise FileNotFoundError(f"Script not found: {script_path}")
+        if not str(script_path).endswith(".js"):
+            raise ValueError(f"Not a JavaScript script: {script_path}")
+
+        full_args = [str(script_path)]
+        if args:
+            full_args.extend(args)
+
+        super().__init__(command=node_cmd, args=full_args, env=env, cwd=cwd)
+        self.script_path = script_path
+
+
+class UvxStdioTransport(StdioTransport):
+    """Transport for running commands via the uvx tool."""
+
+    def __init__(
+        self,
+        tool_name: str,
+        tool_args: list[str] | None = None,
+        project_directory: str | None = None,
+        python_version: str | None = None,
+        with_packages: list[str] | None = None,
+        from_package: str | None = None,
+        env_vars: dict[str, str] | None = None,
+    ):
+        """
+        Initialize a Uvx transport.
+
+        Args:
+            tool_name: Name of the tool to run via uvx
+            tool_args: Arguments to pass to the tool
+            project_directory: Project directory (for package resolution)
+            python_version: Python version to use
+            with_packages: Additional packages to include
+            from_package: Package to install the tool from
+            env_vars: Additional environment variables
+        """
+        # Basic validation
+        if project_directory and not Path(project_directory).exists():
+            raise NotADirectoryError(
+                f"Project directory not found: {project_directory}"
+            )
+
+        # Build uvx arguments
+        uvx_args = []
+        if python_version:
+            uvx_args.extend(["--python", python_version])
+        if from_package:
+            uvx_args.extend(["--from", from_package])
+        for pkg in with_packages or []:
+            uvx_args.extend(["--with", pkg])
+
+        # Add the tool name and tool args
+        uvx_args.append(tool_name)
+        if tool_args:
+            uvx_args.extend(tool_args)
+
+        # Get environment with any additional variables
+        env = None
+        if env_vars:
+            env = os.environ.copy()
+            env.update(env_vars)
+
+        super().__init__(command="uvx", args=uvx_args, env=env, cwd=project_directory)
+        self.tool_name = tool_name
+
+
+class NpxStdioTransport(StdioTransport):
+    """Transport for running commands via the npx tool."""
+
+    def __init__(
+        self,
+        package: str,
+        args: list[str] | None = None,
+        project_directory: str | None = None,
+        env_vars: dict[str, str] | None = None,
+        use_package_lock: bool = True,
+    ):
+        """
+        Initialize an Npx transport.
+
+        Args:
+            package: Name of the npm package to run
+            args: Arguments to pass to the package command
+            project_directory: Project directory with package.json
+            env_vars: Additional environment variables
+            use_package_lock: Whether to use package-lock.json (--prefer-offline)
+        """
+        # Basic validation
+        if project_directory and not Path(project_directory).exists():
+            raise NotADirectoryError(
+                f"Project directory not found: {project_directory}"
+            )
+
+        # Build npx arguments
+        npx_args = []
+        if use_package_lock:
+            npx_args.append("--prefer-offline")
+
+        # Add the package name and args
+        npx_args.append(package)
+        if args:
+            npx_args.extend(args)
+
+        # Get environment with any additional variables
+        env = None
+        if env_vars:
+            env = os.environ.copy()
+            env.update(env_vars)
+
+        super().__init__(command="npx", args=npx_args, env=env, cwd=project_directory)
+        self.package = package
+
+
+class FastMCPTransport(ClientTransport):
+    """
+    Special transport for in-memory connections to an MCP server.
+
+    This is particularly useful for testing or when client and server
+    are in the same process.
+    """
+
+    def __init__(self, mcp: FastMCPServer):
+        self._fastmcp = mcp  # Can be FastMCP or MCPServer
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        # create_connected_server_and_client_session manages the session lifecycle itself
+        async with create_connected_server_and_client_session(
+            server=self._fastmcp._mcp_server,
+            **session_kwargs,
+        ) as session:
+            yield session
+
+    def __repr__(self) -> str:
+        return f"<FastMCP(server='{self._fastmcp.name}')>"
+
+
+def infer_transport(
+    transport: ClientTransport | FastMCPServer | AnyUrl | Path | str,
+) -> ClientTransport:
+    """
+    Infer the appropriate transport type from the given transport argument.
+
+    This function attempts to infer the correct transport type from the provided
+    argument, handling various input types and converting them to the appropriate
+    ClientTransport subclass.
+    """
+    # the transport is already a ClientTransport
+    if isinstance(transport, ClientTransport):
+        return transport
+
+    # the transport is a FastMCP server
+    elif isinstance(transport, FastMCPServer):
+        return FastMCPTransport(mcp=transport)
+
+    # the transport is a path to a script
+    elif isinstance(transport, Path | str) and Path(transport).exists():
+        if str(transport).endswith(".py"):
+            return PythonStdioTransport(script_path=transport)
+        elif str(transport).endswith(".js"):
+            return NodeStdioTransport(script_path=transport)
+        else:
+            raise ValueError(f"Unsupported script type: {transport}")
+
+    # the transport is an http(s) URL
+    elif isinstance(transport, AnyUrl | str) and str(transport).startswith("http"):
+        return SSETransport(url=transport)
+
+    # the transport is a websocket URL
+    elif isinstance(transport, AnyUrl | str) and str(transport).startswith("ws"):
+        return WSTransport(url=transport)
+
+    # the transport is an unknown type
+    else:
+        raise ValueError(f"Could not infer a valid transport from: {transport}")

fastmcp/prompts/__init__.py

@@ -1,4 +1,4 @@
 from .base import Prompt
-from .manager import PromptManager
+from .prompt_manager import PromptManager
 
 __all__ = ["Prompt", "PromptManager"]

fastmcp/prompts/base.py

@@ -1,12 +1,13 @@
 """Base classes for FastMCP prompts."""
 
-import json
-from typing import Any, Callable, Dict, Literal, Optional, Sequence, Awaitable
 import inspect
+import json
+from collections.abc import Awaitable, Callable, Sequence
+from typing import Any, Literal
 
-from pydantic import BaseModel, Field, TypeAdapter, validate_call
-from mcp.types import TextContent, ImageContent, EmbeddedResource
 import pydantic_core
+from mcp.types import EmbeddedResource, ImageContent, TextContent
+from pydantic import BaseModel, Field, TypeAdapter, validate_call
 
 CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource
 
@@ -17,7 +18,7 @@ class Message(BaseModel):
     role: Literal["user", "assistant"]
     content: CONTENT_TYPES
 
-    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
         if isinstance(content, str):
             content = TextContent(type="text", text=content)
         super().__init__(content=content, **kwargs)
@@ -26,22 +27,24 @@ class Message(BaseModel):
 class UserMessage(Message):
     """A message from the user."""
 
-    role: Literal["user"] = "user"
+    role: Literal["user", "assistant"] = "user"
 
-    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
         super().__init__(content=content, **kwargs)
 
 
 class AssistantMessage(Message):
     """A message from the assistant."""
 
-    role: Literal["assistant"] = "assistant"
+    role: Literal["user", "assistant"] = "assistant"
 
-    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
         super().__init__(content=content, **kwargs)
 
 
-message_validator = TypeAdapter(UserMessage | AssistantMessage)
+message_validator = TypeAdapter[UserMessage | AssistantMessage](
+    UserMessage | AssistantMessage
+)
 
 SyncPromptResult = (
     str | Message | dict[str, Any] | Sequence[str | Message | dict[str, Any]]
@@ -71,14 +74,14 @@ class Prompt(BaseModel):
     arguments: list[PromptArgument] | None = Field(
         None, description="Arguments that can be passed to the prompt"
     )
-    fn: Callable = Field(exclude=True)
+    fn: Callable[..., PromptResult | Awaitable[PromptResult]] = Field(exclude=True)
 
     @classmethod
     def from_function(
         cls,
-        fn: Callable[..., PromptResult],
-        name: Optional[str] = None,
-        description: Optional[str] = None,
+        fn: Callable[..., PromptResult | Awaitable[PromptResult]],
+        name: str | None = None,
+        description: str | None = None,
     ) -> "Prompt":
         """Create a Prompt from a function.
 
@@ -97,7 +100,7 @@ class Prompt(BaseModel):
         parameters = TypeAdapter(fn).json_schema()
 
         # Convert parameters to PromptArguments
-        arguments = []
+        arguments: list[PromptArgument] = []
         if "properties" in parameters:
             for param_name, param in parameters["properties"].items():
                 required = param_name in parameters.get("required", [])
@@ -119,7 +122,7 @@ class Prompt(BaseModel):
             fn=fn,
         )
 
-    async def render(self, arguments: Optional[Dict[str, Any]] = None) -> list[Message]:
+    async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]:
         """Render the prompt with arguments."""
         # Validate required arguments
         if self.arguments:
@@ -136,25 +139,23 @@ class Prompt(BaseModel):
                 result = await result
 
             # Validate messages
-            if not isinstance(result, (list, tuple)):
+            if not isinstance(result, list | tuple):
                 result = [result]
 
             # Convert result to messages
-            messages = []
-            for msg in result:
+            messages: list[Message] = []
+            for msg in result:  # type: ignore[reportUnknownVariableType]
                 try:
                     if isinstance(msg, Message):
                         messages.append(msg)
                     elif isinstance(msg, dict):
-                        msg = message_validator.validate_python(msg)
-                        messages.append(msg)
+                        messages.append(message_validator.validate_python(msg))
                     elif isinstance(msg, str):
-                        messages.append(
-                            UserMessage(content=TextContent(type="text", text=msg))
-                        )
+                        content = TextContent(type="text", text=msg)
+                        messages.append(UserMessage(content=content))
                     else:
-                        msg = json.dumps(pydantic_core.to_jsonable_python(msg))
-                        messages.append(Message(role="user", content=msg))
+                        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}"

fastmcp/prompts/prompt_manager.py

@@ -1,9 +1,8 @@
 """Prompt management functionality."""
 
-from typing import Dict, Optional
+from typing import Any
 
-
-from fastmcp.prompts.base import Prompt
+from fastmcp.prompts.base import Message, Prompt
 from fastmcp.utilities.logging import get_logger
 
 logger = get_logger(__name__)
@@ -13,24 +12,63 @@ class PromptManager:
     """Manages FastMCP prompts."""
 
     def __init__(self, warn_on_duplicate_prompts: bool = True):
-        self._prompts: Dict[str, Prompt] = {}
+        self._prompts: dict[str, Prompt] = {}
         self.warn_on_duplicate_prompts = warn_on_duplicate_prompts
 
-    def add_prompt(self, prompt: Prompt) -> Prompt:
+    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."""
-        logger.debug(f"Adding prompt: {prompt.name}")
+
+        # 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
 
-    def get_prompt(self, name: str) -> Optional[Prompt]:
-        """Get prompt by name."""
-        return self._prompts.get(name)
+    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}")
 
-    def list_prompts(self) -> list[Prompt]:
-        """List all registered prompts."""
-        return list(self._prompts.values())
+        return await prompt.render(arguments)
+
+    def import_prompts(
+        self, manager: "PromptManager", prefix: str | None = None
+    ) -> None:
+        """
+        Import all prompts from another PromptManager with prefixed names.
+
+        Args:
+            manager: Another PromptManager instance to import prompts from
+            prefix: Prefix to add to prompt names. The resulting prompt name will
+                   be in the format "{prefix}{original_name}" if prefix is provided,
+                   otherwise the original name is used.
+                   For example, with prefix "weather/" and prompt "forecast_prompt",
+                   the imported prompt would be available as "weather/forecast_prompt"
+        """
+        for name, prompt in manager._prompts.items():
+            # Create prefixed name - we keep the original name in the Prompt object
+            prefixed_name = f"{prefix}{name}" if prefix else name
+
+            # Log the import
+            logger.debug(f"Importing prompt with name {name} as {prefixed_name}")
+
+            # Store the prompt with the prefixed name
+            self._prompts[prefixed_name] = prompt

fastmcp/resources/__init__.py

@@ -1,14 +1,14 @@
 from .base import Resource
+from .resource_manager import ResourceManager
+from .templates import ResourceTemplate
 from .types import (
-    TextResource,
     BinaryResource,
-    FunctionResource,
+    DirectoryResource,
     FileResource,
+    FunctionResource,
     HttpResource,
-    DirectoryResource,
+    TextResource,
 )
-from .templates import ResourceTemplate
-from .resource_manager import ResourceManager
 
 __all__ = [
     "Resource",

fastmcp/resources/base.py

@@ -1,7 +1,7 @@
 """Base classes and interfaces for FastMCP resources."""
 
 import abc
-from typing import Union, Annotated
+from typing import Annotated
 
 from pydantic import (
     AnyUrl,
@@ -43,6 +43,6 @@ class Resource(BaseModel, abc.ABC):
         raise ValueError("Either name or uri must be provided")
 
     @abc.abstractmethod
-    async def read(self) -> Union[str, bytes]:
+    async def read(self) -> str | bytes:
         """Read the resource content."""
         pass