FrameBridge 0.3.1 community source release

This commit is contained in:
FrameBridge Development committed 2026-10-26 22:09:18 -04:00
commit b23819fff1
1354 files changed
+343716

No files matched your search

@@ -0,0 +1,135 @@
from .client.session import ClientSession
from .client.session_group import ClientSessionGroup
from .client.stdio import StdioServerParameters, stdio_client
from .server.session import ServerSession
from .server.stdio import stdio_server
from .shared.exceptions import McpError, UrlElicitationRequiredError
from .types import (
CallToolRequest,
ClientCapabilities,
ClientNotification,
ClientRequest,
ClientResult,
CompleteRequest,
CreateMessageRequest,
CreateMessageResult,
CreateMessageResultWithTools,
ErrorData,
GetPromptRequest,
GetPromptResult,
Implementation,
IncludeContext,
InitializedNotification,
InitializeRequest,
InitializeResult,
JSONRPCError,
JSONRPCRequest,
JSONRPCResponse,
ListPromptsRequest,
ListPromptsResult,
ListResourcesRequest,
ListResourcesResult,
ListToolsResult,
LoggingLevel,
LoggingMessageNotification,
Notification,
PingRequest,
ProgressNotification,
PromptsCapability,
ReadResourceRequest,
ReadResourceResult,
Resource,
ResourcesCapability,
ResourceUpdatedNotification,
RootsCapability,
SamplingCapability,
SamplingContent,
SamplingContextCapability,
SamplingMessage,
SamplingMessageContentBlock,
SamplingToolsCapability,
ServerCapabilities,
ServerNotification,
ServerRequest,
ServerResult,
SetLevelRequest,
StopReason,
SubscribeRequest,
Tool,
ToolChoice,
ToolResultContent,
ToolsCapability,
ToolUseContent,
UnsubscribeRequest,
)
from .types import (
Role as SamplingRole,
)
__all__ = [
"CallToolRequest",
"ClientCapabilities",
"ClientNotification",
"ClientRequest",
"ClientResult",
"ClientSession",
"ClientSessionGroup",
"CompleteRequest",
"CreateMessageRequest",
"CreateMessageResult",
"CreateMessageResultWithTools",
"ErrorData",
"GetPromptRequest",
"GetPromptResult",
"Implementation",
"IncludeContext",
"InitializeRequest",
"InitializeResult",
"InitializedNotification",
"JSONRPCError",
"JSONRPCRequest",
"JSONRPCResponse",
"ListPromptsRequest",
"ListPromptsResult",
"ListResourcesRequest",
"ListResourcesResult",
"ListToolsResult",
"LoggingLevel",
"LoggingMessageNotification",
"McpError",
"Notification",
"PingRequest",
"ProgressNotification",
"PromptsCapability",
"ReadResourceRequest",
"ReadResourceResult",
"Resource",
"ResourcesCapability",
"ResourceUpdatedNotification",
"RootsCapability",
"SamplingCapability",
"SamplingContent",
"SamplingContextCapability",
"SamplingMessage",
"SamplingMessageContentBlock",
"SamplingRole",
"SamplingToolsCapability",
"ServerCapabilities",
"ServerNotification",
"ServerRequest",
"ServerResult",
"ServerSession",
"SetLevelRequest",
"StdioServerParameters",
"StopReason",
"SubscribeRequest",
"Tool",
"ToolChoice",
"ToolResultContent",
"ToolsCapability",
"ToolUseContent",
"UnsubscribeRequest",
"UrlElicitationRequiredError",
"stdio_client",
"stdio_server",
]
@@ -0,0 +1,6 @@
"""FastMCP CLI package."""
from .cli import app
if __name__ == "__main__": # pragma: no cover
app()
@@ -0,0 +1,148 @@
"""Claude app integration utilities."""
import json
import os
import shutil
import sys
from pathlib import Path
from typing import Any
from mcp.server.fastmcp.utilities.logging import get_logger
logger = get_logger(__name__)
MCP_PACKAGE = "mcp[cli]"
def get_claude_config_path() -> Path | None: # pragma: no cover
"""Get the Claude config directory based on platform."""
if sys.platform == "win32":
path = Path(Path.home(), "AppData", "Roaming", "Claude")
elif sys.platform == "darwin":
path = Path(Path.home(), "Library", "Application Support", "Claude")
elif sys.platform.startswith("linux"):
path = Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"), "Claude")
else:
return None
if path.exists():
return path
return None
def get_uv_path() -> str:
"""Get the full path to the uv executable."""
uv_path = shutil.which("uv")
if not uv_path: # pragma: no cover
logger.error(
"uv executable not found in PATH, falling back to 'uv'. Please ensure uv is installed and in your PATH"
)
return "uv" # Fall back to just "uv" if not found
return uv_path
def update_claude_config(
file_spec: str,
server_name: str,
*,
with_editable: Path | None = None,
with_packages: list[str] | None = None,
env_vars: dict[str, str] | None = None,
) -> bool:
"""Add or update a FastMCP server in Claude's configuration.
Args:
file_spec: Path to the server file, optionally with :object suffix
server_name: Name for the server in Claude's config
with_editable: Optional directory to install in editable mode
with_packages: Optional list of additional packages to install
env_vars: Optional dictionary of environment variables. These are merged with
any existing variables, with new values taking precedence.
Raises:
RuntimeError: If Claude Desktop's config directory is not found, indicating
Claude Desktop may not be installed or properly set up.
"""
config_dir = get_claude_config_path()
uv_path = get_uv_path()
if not config_dir: # pragma: no cover
raise RuntimeError(
"Claude Desktop config directory not found. Please ensure Claude Desktop"
" is installed and has been run at least once to initialize its config."
)
config_file = config_dir / "claude_desktop_config.json"
if not config_file.exists(): # pragma: no cover
try:
config_file.write_text("{}")
except Exception:
logger.exception(
"Failed to create Claude config file",
extra={
"config_file": str(config_file),
},
)
return False
try:
config = json.loads(config_file.read_text())
if "mcpServers" not in config:
config["mcpServers"] = {}
# Always preserve existing env vars and merge with new ones
if server_name in config["mcpServers"] and "env" in config["mcpServers"][server_name]: # pragma: no cover
existing_env = config["mcpServers"][server_name]["env"]
if env_vars:
# New vars take precedence over existing ones
env_vars = {**existing_env, **env_vars}
else:
env_vars = existing_env
# Build uv run command
args = ["run", "--frozen"]
# Collect all packages in a set to deduplicate
packages = {MCP_PACKAGE}
if with_packages: # pragma: no cover
packages.update(pkg for pkg in with_packages if pkg)
# Add all packages with --with
for pkg in sorted(packages):
args.extend(["--with", pkg])
if with_editable: # pragma: no cover
args.extend(["--with-editable", str(with_editable)])
# Convert file path to absolute before adding to command
# Split off any :object suffix first
if ":" in file_spec:
file_path, server_object = file_spec.rsplit(":", 1)
file_spec = f"{Path(file_path).resolve()}:{server_object}"
else: # pragma: no cover
file_spec = str(Path(file_spec).resolve())
# Add fastmcp run command
args.extend(["mcp", "run", file_spec])
server_config: dict[str, Any] = {"command": uv_path, "args": args}
# Add environment variables if specified
if env_vars: # pragma: no cover
server_config["env"] = env_vars
config["mcpServers"][server_name] = server_config
config_file.write_text(json.dumps(config, indent=2))
logger.info(
f"Added server '{server_name}' to Claude config",
extra={"config_file": str(config_file)},
)
return True
except Exception: # pragma: no cover
logger.exception(
"Failed to update Claude config",
extra={
"config_file": str(config_file),
},
)
return False
@@ -0,0 +1,488 @@
"""MCP CLI tools."""
import importlib.metadata
import importlib.util
import os
import subprocess
import sys
from pathlib import Path
from typing import Annotated, Any
from mcp.server import FastMCP
from mcp.server import Server as LowLevelServer
try:
import typer
except ImportError: # pragma: no cover
print("Error: typer is required. Install with 'pip install mcp[cli]'")
sys.exit(1)
try:
from mcp.cli import claude
from mcp.server.fastmcp.utilities.logging import get_logger
except ImportError: # pragma: no cover
print("Error: mcp.server.fastmcp is not installed or not in PYTHONPATH")
sys.exit(1)
try:
import dotenv
except ImportError: # pragma: no cover
dotenv = None
logger = get_logger("cli")
app = typer.Typer(
name="mcp",
help="MCP development tools",
add_completion=False,
no_args_is_help=True, # Show help if no args provided
)
def _get_npx_command():
"""Get the correct npx command for the current platform."""
if sys.platform == "win32":
# Try both npx.cmd and npx.exe on Windows
for cmd in ["npx.cmd", "npx.exe", "npx"]:
try:
subprocess.run([cmd, "--version"], check=True, capture_output=True, shell=True)
return cmd
except subprocess.CalledProcessError:
continue
return None
return "npx" # On Unix-like systems, just use npx
def _parse_env_var(env_var: str) -> tuple[str, str]: # pragma: no cover
"""Parse environment variable string in format KEY=VALUE."""
if "=" not in env_var:
logger.error(f"Invalid environment variable format: {env_var}. Must be KEY=VALUE")
sys.exit(1)
key, value = env_var.split("=", 1)
return key.strip(), value.strip()
def _build_uv_command(
file_spec: str,
with_editable: Path | None = None,
with_packages: list[str] | None = None,
) -> list[str]:
"""Build the uv run command that runs an MCP server through mcp run."""
cmd = ["uv"]
cmd.extend(["run", "--with", "mcp"])
if with_editable:
cmd.extend(["--with-editable", str(with_editable)])
if with_packages:
for pkg in with_packages:
if pkg: # pragma: no cover
cmd.extend(["--with", pkg])
# Add mcp run command
cmd.extend(["mcp", "run", file_spec])
return cmd
def _parse_file_path(file_spec: str) -> tuple[Path, str | None]:
"""Parse a file path that may include a server object specification.
Args:
file_spec: Path to file, optionally with :object suffix
Returns:
Tuple of (file_path, server_object)
"""
# First check if we have a Windows path (e.g., C:\...)
has_windows_drive = len(file_spec) > 1 and file_spec[1] == ":"
# Split on the last colon, but only if it's not part of the Windows drive letter
# and there's actually another colon in the string after the drive letter
if ":" in (file_spec[2:] if has_windows_drive else file_spec):
file_str, server_object = file_spec.rsplit(":", 1)
else:
file_str, server_object = file_spec, None
# Resolve the file path
file_path = Path(file_str).expanduser().resolve()
if not file_path.exists():
logger.error(f"File not found: {file_path}")
sys.exit(1)
if not file_path.is_file():
logger.error(f"Not a file: {file_path}")
sys.exit(1)
return file_path, server_object
def _import_server(file: Path, server_object: str | None = None): # pragma: no cover
"""Import an MCP server from a file.
Args:
file: Path to the file
server_object: Optional object name in format "module:object" or just "object"
Returns:
The server object
"""
# Add parent directory to Python path so imports can be resolved
file_dir = str(file.parent)
if file_dir not in sys.path:
sys.path.insert(0, file_dir)
# Import the module
spec = importlib.util.spec_from_file_location("server_module", file)
if not spec or not spec.loader:
logger.error("Could not load module", extra={"file": str(file)})
sys.exit(1)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
def _check_server_object(server_object: Any, object_name: str):
"""Helper function to check that the server object is supported
Args:
server_object: The server object to check.
Returns:
True if it's supported.
"""
if not isinstance(server_object, FastMCP):
logger.error(f"The server object {object_name} is of type {type(server_object)} (expecting {FastMCP}).")
if isinstance(server_object, LowLevelServer):
logger.warning(
"Note that only FastMCP server is supported. Low level Server class is not yet supported."
)
return False
return True
# If no object specified, try common server names
if not server_object:
# Look for the most common server object names
for name in ["mcp", "server", "app"]:
if hasattr(module, name):
if not _check_server_object(getattr(module, name), f"{file}:{name}"):
logger.error(f"Ignoring object '{file}:{name}' as it's not a valid server object")
continue
return getattr(module, name)
logger.error(
f"No server object found in {file}. Please either:\n"
"1. Use a standard variable name (mcp, server, or app)\n"
"2. Specify the object name with file:object syntax"
"3. If the server creates the FastMCP object within main() "
" or another function, refactor the FastMCP object to be a "
" global variable named mcp, server, or app.",
extra={"file": str(file)},
)
sys.exit(1)
# Handle module:object syntax
if ":" in server_object:
module_name, object_name = server_object.split(":", 1)
try:
server_module = importlib.import_module(module_name)
server = getattr(server_module, object_name, None)
except ImportError:
logger.error(
f"Could not import module '{module_name}'",
extra={"file": str(file)},
)
sys.exit(1)
else:
# Just object name
server = getattr(module, server_object, None)
if server is None:
logger.error(
f"Server object '{server_object}' not found",
extra={"file": str(file)},
)
sys.exit(1)
if not _check_server_object(server, server_object):
sys.exit(1)
return server
@app.command()
def version() -> None: # pragma: no cover
"""Show the MCP version."""
try:
version = importlib.metadata.version("mcp")
print(f"MCP version {version}")
except importlib.metadata.PackageNotFoundError:
print("MCP version unknown (package not installed)")
sys.exit(1)
@app.command()
def dev(
file_spec: str = typer.Argument(
...,
help="Python file to run, optionally with :object suffix",
),
with_editable: Annotated[
Path | None,
typer.Option(
"--with-editable",
"-e",
help="Directory containing pyproject.toml to install in editable mode",
exists=True,
file_okay=False,
resolve_path=True,
),
] = None,
with_packages: Annotated[
list[str],
typer.Option(
"--with",
help="Additional packages to install",
),
] = [],
) -> None: # pragma: no cover
"""Run an MCP server with the MCP Inspector."""
file, server_object = _parse_file_path(file_spec)
logger.debug(
"Starting dev server",
extra={
"file": str(file),
"server_object": server_object,
"with_editable": str(with_editable) if with_editable else None,
"with_packages": with_packages,
},
)
try:
# Import server to get dependencies
server = _import_server(file, server_object)
if hasattr(server, "dependencies"):
with_packages = list(set(with_packages + server.dependencies))
uv_cmd = _build_uv_command(file_spec, with_editable, with_packages)
# Get the correct npx command
npx_cmd = _get_npx_command()
if not npx_cmd:
logger.error(
"npx not found. Please ensure Node.js and npm are properly installed and added to your system PATH."
)
sys.exit(1)
# Run the MCP Inspector command with shell=True on Windows
shell = sys.platform == "win32"
process = subprocess.run(
[npx_cmd, "@modelcontextprotocol/inspector"] + uv_cmd,
check=True,
shell=shell,
env=dict(os.environ.items()), # Convert to list of tuples for env update
)
sys.exit(process.returncode)
except subprocess.CalledProcessError as e:
logger.error(
"Dev server failed",
extra={
"file": str(file),
"error": str(e),
"returncode": e.returncode,
},
)
sys.exit(e.returncode)
except FileNotFoundError:
logger.error(
"npx not found. Please ensure Node.js and npm are properly installed "
"and added to your system PATH. You may need to restart your terminal "
"after installation.",
extra={"file": str(file)},
)
sys.exit(1)
@app.command()
def run(
file_spec: str = typer.Argument(
...,
help="Python file to run, optionally with :object suffix",
),
transport: Annotated[
str | None,
typer.Option(
"--transport",
"-t",
help="Transport protocol to use (stdio or sse)",
),
] = None,
) -> None: # pragma: no cover
"""Run an MCP server.
The server can be specified in two ways:\n
1. Module approach: server.py - runs the module directly, expecting a server.run() call.\n
2. Import approach: server.py:app - imports and runs the specified server object.\n\n
Note: This command runs the server directly. You are responsible for ensuring
all dependencies are available.\n
For dependency management, use `mcp install` or `mcp dev` instead.
""" # noqa: E501
file, server_object = _parse_file_path(file_spec)
logger.debug(
"Running server",
extra={
"file": str(file),
"server_object": server_object,
"transport": transport,
},
)
try:
# Import and get server object
server = _import_server(file, server_object)
# Run the server
kwargs = {}
if transport:
kwargs["transport"] = transport
server.run(**kwargs)
except Exception:
logger.exception(
"Failed to run server",
extra={
"file": str(file),
},
)
sys.exit(1)
@app.command()
def install(
file_spec: str = typer.Argument(
...,
help="Python file to run, optionally with :object suffix",
),
server_name: Annotated[
str | None,
typer.Option(
"--name",
"-n",
help="Custom name for the server (defaults to server's name attribute or file name)",
),
] = None,
with_editable: Annotated[
Path | None,
typer.Option(
"--with-editable",
"-e",
help="Directory containing pyproject.toml to install in editable mode",
exists=True,
file_okay=False,
resolve_path=True,
),
] = None,
with_packages: Annotated[
list[str],
typer.Option(
"--with",
help="Additional packages to install",
),
] = [],
env_vars: Annotated[
list[str],
typer.Option(
"--env-var",
"-v",
help="Environment variables in KEY=VALUE format",
),
] = [],
env_file: Annotated[
Path | None,
typer.Option(
"--env-file",
"-f",
help="Load environment variables from a .env file",
exists=True,
file_okay=True,
dir_okay=False,
resolve_path=True,
),
] = None,
) -> None: # pragma: no cover
"""Install an MCP server in the Claude desktop app.
Environment variables are preserved once added and only updated if new values
are explicitly provided.
"""
file, server_object = _parse_file_path(file_spec)
logger.debug(
"Installing server",
extra={
"file": str(file),
"server_name": server_name,
"server_object": server_object,
"with_editable": str(with_editable) if with_editable else None,
"with_packages": with_packages,
},
)
if not claude.get_claude_config_path():
logger.error("Claude app not found")
sys.exit(1)
# Try to import server to get its name, but fall back to file name if dependencies
# missing
name = server_name
server = None
if not name:
try:
server = _import_server(file, server_object)
name = server.name
except (ImportError, ModuleNotFoundError) as e:
logger.debug(
"Could not import server (likely missing dependencies), using file name",
extra={"error": str(e)},
)
name = file.stem
# Get server dependencies if available
server_dependencies = getattr(server, "dependencies", []) if server else []
if server_dependencies:
with_packages = list(set(with_packages + server_dependencies))
# Process environment variables if provided
env_dict: dict[str, str] | None = None
if env_file or env_vars:
env_dict = {}
# Load from .env file if specified
if env_file:
if dotenv:
try:
env_dict |= {k: v for k, v in dotenv.dotenv_values(env_file).items() if v is not None}
except (OSError, ValueError):
logger.exception("Failed to load .env file")
sys.exit(1)
else:
logger.error("python-dotenv is not installed. Cannot load .env file.")
sys.exit(1)
# Add command line environment variables
for env_var in env_vars:
key, value = _parse_env_var(env_var)
env_dict[key] = value
if claude.update_claude_config(
file_spec,
name,
with_editable=with_editable,
with_packages=with_packages,
env_vars=env_dict,
):
logger.info(f"Successfully installed {name} in Claude app")
else:
logger.error(f"Failed to install {name} in Claude app")
sys.exit(1)
Whitespace-only changes.
@@ -0,0 +1,85 @@
import argparse
import logging
import sys
from functools import partial
from urllib.parse import urlparse
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
import mcp.types as types
from mcp.client.session import ClientSession
from mcp.client.sse import sse_client
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
if not sys.warnoptions:
import warnings
warnings.simplefilter("ignore")
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("client")
async def message_handler(
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
) -> None:
if isinstance(message, Exception):
logger.error("Error: %s", message)
return
logger.info("Received message from server: %s", message)
async def run_session(
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
client_info: types.Implementation | None = None,
):
async with ClientSession(
read_stream,
write_stream,
message_handler=message_handler,
client_info=client_info,
) as session:
logger.info("Initializing session")
await session.initialize()
logger.info("Initialized")
async def main(command_or_url: str, args: list[str], env: list[tuple[str, str]]):
env_dict = dict(env)
if urlparse(command_or_url).scheme in ("http", "https"):
# Use SSE client for HTTP(S) URLs
async with sse_client(command_or_url) as streams:
await run_session(*streams)
else:
# Use stdio client for commands
server_parameters = StdioServerParameters(command=command_or_url, args=args, env=env_dict)
async with stdio_client(server_parameters) as streams:
await run_session(*streams)
def cli():
parser = argparse.ArgumentParser()
parser.add_argument("command_or_url", help="Command or URL to connect to")
parser.add_argument("args", nargs="*", help="Additional arguments")
parser.add_argument(
"-e",
"--env",
nargs=2,
action="append",
metavar=("KEY", "VALUE"),
help="Environment variables to set. Can be used multiple times.",
default=[],
)
args = parser.parse_args()
anyio.run(partial(main, args.command_or_url, args.args, args.env), backend="trio")
if __name__ == "__main__":
cli()
@@ -0,0 +1,21 @@
"""
OAuth2 Authentication implementation for HTTPX.
Implements authorization code flow with PKCE and automatic token refresh.
"""
from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError
from mcp.client.auth.oauth2 import (
OAuthClientProvider,
PKCEParameters,
TokenStorage,
)
__all__ = [
"OAuthClientProvider",
"OAuthFlowError",
"OAuthRegistrationError",
"OAuthTokenError",
"PKCEParameters",
"TokenStorage",
]
@@ -0,0 +1,10 @@
class OAuthFlowError(Exception):
"""Base exception for OAuth flow errors."""
class OAuthTokenError(OAuthFlowError):
"""Raised when token operations fail."""
class OAuthRegistrationError(OAuthFlowError):
"""Raised when client registration fails."""
@@ -0,0 +1,487 @@
"""
OAuth client credential extensions for MCP.
Provides OAuth providers for machine-to-machine authentication flows:
- ClientCredentialsOAuthProvider: For client_credentials with client_id + client_secret
- PrivateKeyJWTOAuthProvider: For client_credentials with private_key_jwt authentication
(typically using a pre-built JWT from workload identity federation)
- RFC7523OAuthClientProvider: For jwt-bearer grant (RFC 7523 Section 2.1)
"""
import time
from collections.abc import Awaitable, Callable
from typing import Any, Literal
from uuid import uuid4
import httpx
import jwt
from pydantic import BaseModel, Field
from mcp.client.auth import OAuthClientProvider, OAuthFlowError, OAuthTokenError, TokenStorage
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
class ClientCredentialsOAuthProvider(OAuthClientProvider):
"""OAuth provider for client_credentials grant with client_id + client_secret.
This provider sets client_info directly, bypassing dynamic client registration.
Use this when you already have client credentials (client_id and client_secret).
Example:
```python
provider = ClientCredentialsOAuthProvider(
server_url="https://api.example.com",
storage=my_token_storage,
client_id="my-client-id",
client_secret="my-client-secret",
)
```
"""
def __init__(
self,
server_url: str,
storage: TokenStorage,
client_id: str,
client_secret: str,
token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] = "client_secret_basic",
scopes: str | None = None,
) -> None:
"""Initialize client_credentials OAuth provider.
Args:
server_url: The MCP server URL.
storage: Token storage implementation.
client_id: The OAuth client ID.
client_secret: The OAuth client secret.
token_endpoint_auth_method: Authentication method for token endpoint.
Either "client_secret_basic" (default) or "client_secret_post".
scopes: Optional space-separated list of scopes to request.
"""
# Build minimal client_metadata for the base class
client_metadata = OAuthClientMetadata(
redirect_uris=None,
grant_types=["client_credentials"],
token_endpoint_auth_method=token_endpoint_auth_method,
scope=scopes,
)
super().__init__(server_url, client_metadata, storage, None, None, 300.0)
# Store client_info to be set during _initialize - no dynamic registration needed
self._fixed_client_info = OAuthClientInformationFull(
redirect_uris=None,
client_id=client_id,
client_secret=client_secret,
grant_types=["client_credentials"],
token_endpoint_auth_method=token_endpoint_auth_method,
scope=scopes,
)
async def _initialize(self) -> None:
"""Load stored tokens and set pre-configured client_info."""
self.context.current_tokens = await self.context.storage.get_tokens()
self.context.client_info = self._fixed_client_info
self._initialized = True
async def _perform_authorization(self) -> httpx.Request:
"""Perform client_credentials authorization."""
return await self._exchange_token_client_credentials()
async def _exchange_token_client_credentials(self) -> httpx.Request:
"""Build token exchange request for client_credentials grant."""
token_data: dict[str, Any] = {
"grant_type": "client_credentials",
}
headers: dict[str, str] = {"Content-Type": "application/x-www-form-urlencoded"}
# Use standard auth methods (client_secret_basic, client_secret_post, none)
token_data, headers = self.context.prepare_token_auth(token_data, headers)
if self.context.should_include_resource_param(self.context.protocol_version):
token_data["resource"] = self.context.get_resource_url()
if self.context.client_metadata.scope:
token_data["scope"] = self.context.client_metadata.scope
token_url = self._get_token_endpoint()
return httpx.Request("POST", token_url, data=token_data, headers=headers)
def static_assertion_provider(token: str) -> Callable[[str], Awaitable[str]]:
"""Create an assertion provider that returns a static JWT token.
Use this when you have a pre-built JWT (e.g., from workload identity federation)
that doesn't need the audience parameter.
Example:
```python
provider = PrivateKeyJWTOAuthProvider(
server_url="https://api.example.com",
storage=my_token_storage,
client_id="my-client-id",
assertion_provider=static_assertion_provider(my_prebuilt_jwt),
)
```
Args:
token: The pre-built JWT assertion string.
Returns:
An async callback suitable for use as an assertion_provider.
"""
async def provider(audience: str) -> str:
return token
return provider
class SignedJWTParameters(BaseModel):
"""Parameters for creating SDK-signed JWT assertions.
Use `create_assertion_provider()` to create an assertion provider callback
for use with `PrivateKeyJWTOAuthProvider`.
Example:
```python
jwt_params = SignedJWTParameters(
issuer="my-client-id",
subject="my-client-id",
signing_key=private_key_pem,
)
provider = PrivateKeyJWTOAuthProvider(
server_url="https://api.example.com",
storage=my_token_storage,
client_id="my-client-id",
assertion_provider=jwt_params.create_assertion_provider(),
)
```
"""
issuer: str = Field(description="Issuer for JWT assertions (typically client_id).")
subject: str = Field(description="Subject identifier for JWT assertions (typically client_id).")
signing_key: str = Field(description="Private key for JWT signing (PEM format).")
signing_algorithm: str = Field(default="RS256", description="Algorithm for signing JWT assertions.")
lifetime_seconds: int = Field(default=300, description="Lifetime of generated JWT in seconds.")
additional_claims: dict[str, Any] | None = Field(default=None, description="Additional claims.")
def create_assertion_provider(self) -> Callable[[str], Awaitable[str]]:
"""Create an assertion provider callback for use with PrivateKeyJWTOAuthProvider.
Returns:
An async callback that takes the audience (authorization server issuer URL)
and returns a signed JWT assertion.
"""
async def provider(audience: str) -> str:
now = int(time.time())
claims: dict[str, Any] = {
"iss": self.issuer,
"sub": self.subject,
"aud": audience,
"exp": now + self.lifetime_seconds,
"iat": now,
"jti": str(uuid4()),
}
if self.additional_claims:
claims.update(self.additional_claims)
return jwt.encode(claims, self.signing_key, algorithm=self.signing_algorithm)
return provider
class PrivateKeyJWTOAuthProvider(OAuthClientProvider):
"""OAuth provider for client_credentials grant with private_key_jwt authentication.
Uses RFC 7523 Section 2.2 for client authentication via JWT assertion.
The JWT assertion's audience MUST be the authorization server's issuer identifier
(per RFC 7523bis security updates). The `assertion_provider` callback receives
this audience value and must return a JWT with that audience.
**Option 1: Pre-built JWT via Workload Identity Federation**
In production scenarios, the JWT assertion is typically obtained from a workload
identity provider (e.g., GCP, AWS IAM, Azure AD):
```python
async def get_workload_identity_token(audience: str) -> str:
# Fetch JWT from your identity provider
# The JWT's audience must match the provided audience parameter
return await fetch_token_from_identity_provider(audience=audience)
provider = PrivateKeyJWTOAuthProvider(
server_url="https://api.example.com",
storage=my_token_storage,
client_id="my-client-id",
assertion_provider=get_workload_identity_token,
)
```
**Option 2: Static pre-built JWT**
If you have a static JWT that doesn't need the audience parameter:
```python
provider = PrivateKeyJWTOAuthProvider(
server_url="https://api.example.com",
storage=my_token_storage,
client_id="my-client-id",
assertion_provider=static_assertion_provider(my_prebuilt_jwt),
)
```
**Option 3: SDK-signed JWT (for testing/simple setups)**
For testing or simple deployments, use `SignedJWTParameters.create_assertion_provider()`:
```python
jwt_params = SignedJWTParameters(
issuer="my-client-id",
subject="my-client-id",
signing_key=private_key_pem,
)
provider = PrivateKeyJWTOAuthProvider(
server_url="https://api.example.com",
storage=my_token_storage,
client_id="my-client-id",
assertion_provider=jwt_params.create_assertion_provider(),
)
```
"""
def __init__(
self,
server_url: str,
storage: TokenStorage,
client_id: str,
assertion_provider: Callable[[str], Awaitable[str]],
scopes: str | None = None,
) -> None:
"""Initialize private_key_jwt OAuth provider.
Args:
server_url: The MCP server URL.
storage: Token storage implementation.
client_id: The OAuth client ID.
assertion_provider: Async callback that takes the audience (authorization
server's issuer identifier) and returns a JWT assertion. Use
`SignedJWTParameters.create_assertion_provider()` for SDK-signed JWTs,
`static_assertion_provider()` for pre-built JWTs, or provide your own
callback for workload identity federation.
scopes: Optional space-separated list of scopes to request.
"""
# Build minimal client_metadata for the base class
client_metadata = OAuthClientMetadata(
redirect_uris=None,
grant_types=["client_credentials"],
token_endpoint_auth_method="private_key_jwt",
scope=scopes,
)
super().__init__(server_url, client_metadata, storage, None, None, 300.0)
self._assertion_provider = assertion_provider
# Store client_info to be set during _initialize - no dynamic registration needed
self._fixed_client_info = OAuthClientInformationFull(
redirect_uris=None,
client_id=client_id,
grant_types=["client_credentials"],
token_endpoint_auth_method="private_key_jwt",
scope=scopes,
)
async def _initialize(self) -> None:
"""Load stored tokens and set pre-configured client_info."""
self.context.current_tokens = await self.context.storage.get_tokens()
self.context.client_info = self._fixed_client_info
self._initialized = True
async def _perform_authorization(self) -> httpx.Request:
"""Perform client_credentials authorization with private_key_jwt."""
return await self._exchange_token_client_credentials()
async def _add_client_authentication_jwt(self, *, token_data: dict[str, Any]) -> None:
"""Add JWT assertion for client authentication to token endpoint parameters."""
if not self.context.oauth_metadata:
raise OAuthFlowError("Missing OAuth metadata for private_key_jwt flow") # pragma: no cover
# Audience MUST be the issuer identifier of the authorization server
# https://datatracker.ietf.org/doc/html/draft-ietf-oauth-rfc7523bis-01
audience = str(self.context.oauth_metadata.issuer)
assertion = await self._assertion_provider(audience)
# RFC 7523 Section 2.2: client authentication via JWT
token_data["client_assertion"] = assertion
token_data["client_assertion_type"] = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
async def _exchange_token_client_credentials(self) -> httpx.Request:
"""Build token exchange request for client_credentials grant with private_key_jwt."""
token_data: dict[str, Any] = {
"grant_type": "client_credentials",
}
headers: dict[str, str] = {"Content-Type": "application/x-www-form-urlencoded"}
# Add JWT client authentication (RFC 7523 Section 2.2)
await self._add_client_authentication_jwt(token_data=token_data)
if self.context.should_include_resource_param(self.context.protocol_version):
token_data["resource"] = self.context.get_resource_url()
if self.context.client_metadata.scope:
token_data["scope"] = self.context.client_metadata.scope
token_url = self._get_token_endpoint()
return httpx.Request("POST", token_url, data=token_data, headers=headers)
class JWTParameters(BaseModel):
"""JWT parameters."""
assertion: str | None = Field(
default=None,
description="JWT assertion for JWT authentication. "
"Will be used instead of generating a new assertion if provided.",
)
issuer: str | None = Field(default=None, description="Issuer for JWT assertions.")
subject: str | None = Field(default=None, description="Subject identifier for JWT assertions.")
audience: str | None = Field(default=None, description="Audience for JWT assertions.")
claims: dict[str, Any] | None = Field(default=None, description="Additional claims for JWT assertions.")
jwt_signing_algorithm: str | None = Field(default="RS256", description="Algorithm for signing JWT assertions.")
jwt_signing_key: str | None = Field(default=None, description="Private key for JWT signing.")
jwt_lifetime_seconds: int = Field(default=300, description="Lifetime of generated JWT in seconds.")
def to_assertion(self, with_audience_fallback: str | None = None) -> str:
if self.assertion is not None:
# Prebuilt JWT (e.g. acquired out-of-band)
assertion = self.assertion
else:
if not self.jwt_signing_key:
raise OAuthFlowError("Missing signing key for JWT bearer grant") # pragma: no cover
if not self.issuer:
raise OAuthFlowError("Missing issuer for JWT bearer grant") # pragma: no cover
if not self.subject:
raise OAuthFlowError("Missing subject for JWT bearer grant") # pragma: no cover
audience = self.audience if self.audience else with_audience_fallback
if not audience:
raise OAuthFlowError("Missing audience for JWT bearer grant") # pragma: no cover
now = int(time.time())
claims: dict[str, Any] = {
"iss": self.issuer,
"sub": self.subject,
"aud": audience,
"exp": now + self.jwt_lifetime_seconds,
"iat": now,
"jti": str(uuid4()),
}
claims.update(self.claims or {})
assertion = jwt.encode(
claims,
self.jwt_signing_key,
algorithm=self.jwt_signing_algorithm or "RS256",
)
return assertion
class RFC7523OAuthClientProvider(OAuthClientProvider):
"""OAuth client provider for RFC 7523 jwt-bearer grant.
.. deprecated::
Use :class:`ClientCredentialsOAuthProvider` for client_credentials with
client_id + client_secret, or :class:`PrivateKeyJWTOAuthProvider` for
client_credentials with private_key_jwt authentication instead.
This provider supports the jwt-bearer authorization grant (RFC 7523 Section 2.1)
where the JWT itself is the authorization grant.
"""
def __init__(
self,
server_url: str,
client_metadata: OAuthClientMetadata,
storage: TokenStorage,
redirect_handler: Callable[[str], Awaitable[None]] | None = None,
callback_handler: Callable[[], Awaitable[tuple[str, str | None]]] | None = None,
timeout: float = 300.0,
jwt_parameters: JWTParameters | None = None,
) -> None:
import warnings
warnings.warn(
"RFC7523OAuthClientProvider is deprecated. Use ClientCredentialsOAuthProvider "
"or PrivateKeyJWTOAuthProvider instead.",
DeprecationWarning,
stacklevel=2,
)
super().__init__(server_url, client_metadata, storage, redirect_handler, callback_handler, timeout)
self.jwt_parameters = jwt_parameters
async def _exchange_token_authorization_code(
self, auth_code: str, code_verifier: str, *, token_data: dict[str, Any] | None = None
) -> httpx.Request: # pragma: no cover
"""Build token exchange request for authorization_code flow."""
token_data = token_data or {}
if self.context.client_metadata.token_endpoint_auth_method == "private_key_jwt":
self._add_client_authentication_jwt(token_data=token_data)
return await super()._exchange_token_authorization_code(auth_code, code_verifier, token_data=token_data)
async def _perform_authorization(self) -> httpx.Request: # pragma: no cover
"""Perform the authorization flow."""
if "urn:ietf:params:oauth:grant-type:jwt-bearer" in self.context.client_metadata.grant_types:
token_request = await self._exchange_token_jwt_bearer()
return token_request
else:
return await super()._perform_authorization()
def _add_client_authentication_jwt(self, *, token_data: dict[str, Any]): # pragma: no cover
"""Add JWT assertion for client authentication to token endpoint parameters."""
if not self.jwt_parameters:
raise OAuthTokenError("Missing JWT parameters for private_key_jwt flow")
if not self.context.oauth_metadata:
raise OAuthTokenError("Missing OAuth metadata for private_key_jwt flow")
# We need to set the audience to the issuer identifier of the authorization server
# https://datatracker.ietf.org/doc/html/draft-ietf-oauth-rfc7523bis-01#name-updates-to-rfc-7523
issuer = str(self.context.oauth_metadata.issuer)
assertion = self.jwt_parameters.to_assertion(with_audience_fallback=issuer)
# When using private_key_jwt, in a client_credentials flow, we use RFC 7523 Section 2.2
token_data["client_assertion"] = assertion
token_data["client_assertion_type"] = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
# We need to set the audience to the resource server, the audience is difference from the one in claims
# it represents the resource server that will validate the token
token_data["audience"] = self.context.get_resource_url()
async def _exchange_token_jwt_bearer(self) -> httpx.Request:
"""Build token exchange request for JWT bearer grant."""
if not self.context.client_info:
raise OAuthFlowError("Missing client info") # pragma: no cover
if not self.jwt_parameters:
raise OAuthFlowError("Missing JWT parameters") # pragma: no cover
if not self.context.oauth_metadata:
raise OAuthTokenError("Missing OAuth metadata") # pragma: no cover
# We need to set the audience to the issuer identifier of the authorization server
# https://datatracker.ietf.org/doc/html/draft-ietf-oauth-rfc7523bis-01#name-updates-to-rfc-7523
issuer = str(self.context.oauth_metadata.issuer)
assertion = self.jwt_parameters.to_assertion(with_audience_fallback=issuer)
token_data = {
"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
"assertion": assertion,
}
if self.context.should_include_resource_param(self.context.protocol_version): # pragma: no branch
token_data["resource"] = self.context.get_resource_url()
if self.context.client_metadata.scope: # pragma: no branch
token_data["scope"] = self.context.client_metadata.scope
token_url = self._get_token_endpoint()
return httpx.Request(
"POST", token_url, data=token_data, headers={"Content-Type": "application/x-www-form-urlencoded"}
)
@@ -0,0 +1,616 @@
"""
OAuth2 Authentication implementation for HTTPX.
Implements authorization code flow with PKCE and automatic token refresh.
"""
import base64
import hashlib
import logging
import secrets
import string
import time
from collections.abc import AsyncGenerator, Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any, Protocol
from urllib.parse import quote, urlencode, urljoin, urlparse
import anyio
import httpx
from pydantic import BaseModel, Field, ValidationError
from mcp.client.auth.exceptions import OAuthFlowError, OAuthTokenError
from mcp.client.auth.utils import (
build_oauth_authorization_server_metadata_discovery_urls,
build_protected_resource_metadata_discovery_urls,
create_client_info_from_metadata_url,
create_client_registration_request,
create_oauth_metadata_request,
extract_field_from_www_auth,
extract_resource_metadata_from_www_auth,
extract_scope_from_www_auth,
get_client_metadata_scopes,
handle_auth_metadata_response,
handle_protected_resource_response,
handle_registration_response,
handle_token_response_scopes,
is_valid_client_metadata_url,
should_use_client_metadata_url,
)
from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
from mcp.shared.auth import (
OAuthClientInformationFull,
OAuthClientMetadata,
OAuthMetadata,
OAuthToken,
ProtectedResourceMetadata,
)
from mcp.shared.auth_utils import (
calculate_token_expiry,
check_resource_allowed,
resource_url_from_server_url,
)
logger = logging.getLogger(__name__)
class PKCEParameters(BaseModel):
"""PKCE (Proof Key for Code Exchange) parameters."""
code_verifier: str = Field(..., min_length=43, max_length=128)
code_challenge: str = Field(..., min_length=43, max_length=128)
@classmethod
def generate(cls) -> "PKCEParameters":
"""Generate new PKCE parameters."""
code_verifier = "".join(secrets.choice(string.ascii_letters + string.digits + "-._~") for _ in range(128))
digest = hashlib.sha256(code_verifier.encode()).digest()
code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=")
return cls(code_verifier=code_verifier, code_challenge=code_challenge)
class TokenStorage(Protocol):
"""Protocol for token storage implementations."""
async def get_tokens(self) -> OAuthToken | None:
"""Get stored tokens."""
...
async def set_tokens(self, tokens: OAuthToken) -> None:
"""Store tokens."""
...
async def get_client_info(self) -> OAuthClientInformationFull | None:
"""Get stored client information."""
...
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
"""Store client information."""
...
@dataclass
class OAuthContext:
"""OAuth flow context."""
server_url: str
client_metadata: OAuthClientMetadata
storage: TokenStorage
redirect_handler: Callable[[str], Awaitable[None]] | None
callback_handler: Callable[[], Awaitable[tuple[str, str | None]]] | None
timeout: float = 300.0
client_metadata_url: str | None = None
# Discovered metadata
protected_resource_metadata: ProtectedResourceMetadata | None = None
oauth_metadata: OAuthMetadata | None = None
auth_server_url: str | None = None
protocol_version: str | None = None
# Client registration
client_info: OAuthClientInformationFull | None = None
# Token management
current_tokens: OAuthToken | None = None
token_expiry_time: float | None = None
# State
lock: anyio.Lock = field(default_factory=anyio.Lock)
def get_authorization_base_url(self, server_url: str) -> str:
"""Extract base URL by removing path component."""
parsed = urlparse(server_url)
return f"{parsed.scheme}://{parsed.netloc}"
def update_token_expiry(self, token: OAuthToken) -> None:
"""Update token expiry time using shared util function."""
self.token_expiry_time = calculate_token_expiry(token.expires_in)
def is_token_valid(self) -> bool:
"""Check if current token is valid."""
return bool(
self.current_tokens
and self.current_tokens.access_token
and (not self.token_expiry_time or time.time() <= self.token_expiry_time)
)
def can_refresh_token(self) -> bool:
"""Check if token can be refreshed."""
return bool(self.current_tokens and self.current_tokens.refresh_token and self.client_info)
def clear_tokens(self) -> None:
"""Clear current tokens."""
self.current_tokens = None
self.token_expiry_time = None
def get_resource_url(self) -> str:
"""Get resource URL for RFC 8707.
Uses PRM resource if it's a valid parent, otherwise uses canonical server URL.
"""
resource = resource_url_from_server_url(self.server_url)
# If PRM provides a resource that's a valid parent, use it
if self.protected_resource_metadata and self.protected_resource_metadata.resource:
prm_resource = str(self.protected_resource_metadata.resource)
if check_resource_allowed(requested_resource=resource, configured_resource=prm_resource):
resource = prm_resource
return resource
def should_include_resource_param(self, protocol_version: str | None = None) -> bool:
"""Determine if the resource parameter should be included in OAuth requests.
Returns True if:
- Protected resource metadata is available, OR
- MCP-Protocol-Version header is 2025-06-18 or later
"""
# If we have protected resource metadata, include the resource param
if self.protected_resource_metadata is not None:
return True
# If no protocol version provided, don't include resource param
if not protocol_version:
return False
# Check if protocol version is 2025-06-18 or later
# Version format is YYYY-MM-DD, so string comparison works
return protocol_version >= "2025-06-18"
def prepare_token_auth(
self, data: dict[str, str], headers: dict[str, str] | None = None
) -> tuple[dict[str, str], dict[str, str]]:
"""Prepare authentication for token requests.
Args:
data: The form data to send
headers: Optional headers dict to update
Returns:
Tuple of (updated_data, updated_headers)
"""
if headers is None:
headers = {} # pragma: no cover
if not self.client_info:
return data, headers # pragma: no cover
auth_method = self.client_info.token_endpoint_auth_method
if auth_method == "client_secret_basic" and self.client_info.client_id and self.client_info.client_secret:
# URL-encode client ID and secret per RFC 6749 Section 2.3.1
encoded_id = quote(self.client_info.client_id, safe="")
encoded_secret = quote(self.client_info.client_secret, safe="")
credentials = f"{encoded_id}:{encoded_secret}"
encoded_credentials = base64.b64encode(credentials.encode()).decode()
headers["Authorization"] = f"Basic {encoded_credentials}"
# Don't include client_secret in body for basic auth
data = {k: v for k, v in data.items() if k != "client_secret"}
elif auth_method == "client_secret_post" and self.client_info.client_secret:
# Include client_secret in request body
data["client_secret"] = self.client_info.client_secret
# For auth_method == "none", don't add any client_secret
return data, headers
class OAuthClientProvider(httpx.Auth):
"""
OAuth2 authentication for httpx.
Handles OAuth flow with automatic client registration and token storage.
"""
requires_response_body = True
def __init__(
self,
server_url: str,
client_metadata: OAuthClientMetadata,
storage: TokenStorage,
redirect_handler: Callable[[str], Awaitable[None]] | None = None,
callback_handler: Callable[[], Awaitable[tuple[str, str | None]]] | None = None,
timeout: float = 300.0,
client_metadata_url: str | None = None,
):
"""Initialize OAuth2 authentication.
Args:
server_url: The MCP server URL.
client_metadata: OAuth client metadata for registration.
storage: Token storage implementation.
redirect_handler: Handler for authorization redirects.
callback_handler: Handler for authorization callbacks.
timeout: Timeout for the OAuth flow.
client_metadata_url: URL-based client ID. When provided and the server
advertises client_id_metadata_document_supported=true, this URL will be
used as the client_id instead of performing dynamic client registration.
Must be a valid HTTPS URL with a non-root pathname.
Raises:
ValueError: If client_metadata_url is provided but not a valid HTTPS URL
with a non-root pathname.
"""
# Validate client_metadata_url if provided
if client_metadata_url is not None and not is_valid_client_metadata_url(client_metadata_url):
raise ValueError(
f"client_metadata_url must be a valid HTTPS URL with a non-root pathname, got: {client_metadata_url}"
)
self.context = OAuthContext(
server_url=server_url,
client_metadata=client_metadata,
storage=storage,
redirect_handler=redirect_handler,
callback_handler=callback_handler,
timeout=timeout,
client_metadata_url=client_metadata_url,
)
self._initialized = False
async def _handle_protected_resource_response(self, response: httpx.Response) -> bool:
"""
Handle protected resource metadata discovery response.
Per SEP-985, supports fallback when discovery fails at one URL.
Returns:
True if metadata was successfully discovered, False if we should try next URL
"""
if response.status_code == 200:
try:
content = await response.aread()
metadata = ProtectedResourceMetadata.model_validate_json(content)
self.context.protected_resource_metadata = metadata
if metadata.authorization_servers: # pragma: no branch
self.context.auth_server_url = str(metadata.authorization_servers[0])
return True
except ValidationError: # pragma: no cover
# Invalid metadata - try next URL
logger.warning(f"Invalid protected resource metadata at {response.request.url}")
return False
elif response.status_code == 404: # pragma: no cover
# Not found - try next URL in fallback chain
logger.debug(f"Protected resource metadata not found at {response.request.url}, trying next URL")
return False
else:
# Other error - fail immediately
raise OAuthFlowError(
f"Protected Resource Metadata request failed: {response.status_code}"
) # pragma: no cover
async def _perform_authorization(self) -> httpx.Request:
"""Perform the authorization flow."""
auth_code, code_verifier = await self._perform_authorization_code_grant()
token_request = await self._exchange_token_authorization_code(auth_code, code_verifier)
return token_request
async def _perform_authorization_code_grant(self) -> tuple[str, str]:
"""Perform the authorization redirect and get auth code."""
if self.context.client_metadata.redirect_uris is None:
raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover
if not self.context.redirect_handler:
raise OAuthFlowError("No redirect handler provided for authorization code grant") # pragma: no cover
if not self.context.callback_handler:
raise OAuthFlowError("No callback handler provided for authorization code grant") # pragma: no cover
if self.context.oauth_metadata and self.context.oauth_metadata.authorization_endpoint:
auth_endpoint = str(self.context.oauth_metadata.authorization_endpoint) # pragma: no cover
else:
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
auth_endpoint = urljoin(auth_base_url, "/authorize")
if not self.context.client_info:
raise OAuthFlowError("No client info available for authorization") # pragma: no cover
# Generate PKCE parameters
pkce_params = PKCEParameters.generate()
state = secrets.token_urlsafe(32)
auth_params = {
"response_type": "code",
"client_id": self.context.client_info.client_id,
"redirect_uri": str(self.context.client_metadata.redirect_uris[0]),
"state": state,
"code_challenge": pkce_params.code_challenge,
"code_challenge_method": "S256",
}
# Only include resource param if conditions are met
if self.context.should_include_resource_param(self.context.protocol_version):
auth_params["resource"] = self.context.get_resource_url() # RFC 8707 # pragma: no cover
if self.context.client_metadata.scope: # pragma: no branch
auth_params["scope"] = self.context.client_metadata.scope
authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}"
await self.context.redirect_handler(authorization_url)
# Wait for callback
auth_code, returned_state = await self.context.callback_handler()
if returned_state is None or not secrets.compare_digest(returned_state, state):
raise OAuthFlowError(f"State parameter mismatch: {returned_state} != {state}") # pragma: no cover
if not auth_code:
raise OAuthFlowError("No authorization code received") # pragma: no cover
# Return auth code and code verifier for token exchange
return auth_code, pkce_params.code_verifier
def _get_token_endpoint(self) -> str:
if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint:
token_url = str(self.context.oauth_metadata.token_endpoint)
else:
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
token_url = urljoin(auth_base_url, "/token")
return token_url
async def _exchange_token_authorization_code(
self, auth_code: str, code_verifier: str, *, token_data: dict[str, Any] | None = {}
) -> httpx.Request:
"""Build token exchange request for authorization_code flow."""
if self.context.client_metadata.redirect_uris is None:
raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover
if not self.context.client_info:
raise OAuthFlowError("Missing client info") # pragma: no cover
token_url = self._get_token_endpoint()
token_data = token_data or {}
token_data.update(
{
"grant_type": "authorization_code",
"code": auth_code,
"redirect_uri": str(self.context.client_metadata.redirect_uris[0]),
"client_id": self.context.client_info.client_id,
"code_verifier": code_verifier,
}
)
# Only include resource param if conditions are met
if self.context.should_include_resource_param(self.context.protocol_version):
token_data["resource"] = self.context.get_resource_url() # RFC 8707
# Prepare authentication based on preferred method
headers = {"Content-Type": "application/x-www-form-urlencoded"}
token_data, headers = self.context.prepare_token_auth(token_data, headers)
return httpx.Request("POST", token_url, data=token_data, headers=headers)
async def _handle_token_response(self, response: httpx.Response) -> None:
"""Handle token exchange response."""
if response.status_code != 200:
body = await response.aread() # pragma: no cover
body_text = body.decode("utf-8") # pragma: no cover
raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body_text}") # pragma: no cover
# Parse and validate response with scope validation
token_response = await handle_token_response_scopes(response)
# Store tokens in context
self.context.current_tokens = token_response
self.context.update_token_expiry(token_response)
await self.context.storage.set_tokens(token_response)
async def _refresh_token(self) -> httpx.Request:
"""Build token refresh request."""
if not self.context.current_tokens or not self.context.current_tokens.refresh_token:
raise OAuthTokenError("No refresh token available") # pragma: no cover
if not self.context.client_info or not self.context.client_info.client_id:
raise OAuthTokenError("No client info available") # pragma: no cover
if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint:
token_url = str(self.context.oauth_metadata.token_endpoint) # pragma: no cover
else:
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
token_url = urljoin(auth_base_url, "/token")
refresh_data: dict[str, str] = {
"grant_type": "refresh_token",
"refresh_token": self.context.current_tokens.refresh_token,
"client_id": self.context.client_info.client_id,
}
# Only include resource param if conditions are met
if self.context.should_include_resource_param(self.context.protocol_version):
refresh_data["resource"] = self.context.get_resource_url() # RFC 8707
# Prepare authentication based on preferred method
headers = {"Content-Type": "application/x-www-form-urlencoded"}
refresh_data, headers = self.context.prepare_token_auth(refresh_data, headers)
return httpx.Request("POST", token_url, data=refresh_data, headers=headers)
async def _handle_refresh_response(self, response: httpx.Response) -> bool: # pragma: no cover
"""Handle token refresh response. Returns True if successful."""
if response.status_code != 200:
logger.warning(f"Token refresh failed: {response.status_code}")
self.context.clear_tokens()
return False
try:
content = await response.aread()
token_response = OAuthToken.model_validate_json(content)
self.context.current_tokens = token_response
self.context.update_token_expiry(token_response)
await self.context.storage.set_tokens(token_response)
return True
except ValidationError:
logger.exception("Invalid refresh response")
self.context.clear_tokens()
return False
async def _initialize(self) -> None: # pragma: no cover
"""Load stored tokens and client info."""
self.context.current_tokens = await self.context.storage.get_tokens()
self.context.client_info = await self.context.storage.get_client_info()
self._initialized = True
def _add_auth_header(self, request: httpx.Request) -> None:
"""Add authorization header to request if we have valid tokens."""
if self.context.current_tokens and self.context.current_tokens.access_token: # pragma: no branch
request.headers["Authorization"] = f"Bearer {self.context.current_tokens.access_token}"
async def _handle_oauth_metadata_response(self, response: httpx.Response) -> None:
content = await response.aread()
metadata = OAuthMetadata.model_validate_json(content)
self.context.oauth_metadata = metadata
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
"""HTTPX auth flow integration."""
async with self.context.lock:
if not self._initialized:
await self._initialize() # pragma: no cover
# Capture protocol version from request headers
self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION)
if not self.context.is_token_valid() and self.context.can_refresh_token():
# Try to refresh token
refresh_request = await self._refresh_token() # pragma: no cover
refresh_response = yield refresh_request # pragma: no cover
if not await self._handle_refresh_response(refresh_response): # pragma: no cover
# Refresh failed, need full re-authentication
self._initialized = False
if self.context.is_token_valid():
self._add_auth_header(request)
response = yield request
if response.status_code == 401:
# Perform full OAuth flow
try:
# OAuth flow must be inline due to generator constraints
www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response)
# Step 1: Discover protected resource metadata (SEP-985 with fallback support)
prm_discovery_urls = build_protected_resource_metadata_discovery_urls(
www_auth_resource_metadata_url, self.context.server_url
)
for url in prm_discovery_urls: # pragma: no branch
discovery_request = create_oauth_metadata_request(url)
discovery_response = yield discovery_request # sending request
prm = await handle_protected_resource_response(discovery_response)
if prm:
self.context.protected_resource_metadata = prm
# todo: try all authorization_servers to find the OASM
assert (
len(prm.authorization_servers) > 0
) # this is always true as authorization_servers has a min length of 1
self.context.auth_server_url = str(prm.authorization_servers[0])
break
else:
logger.debug(f"Protected resource metadata discovery failed: {url}")
asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls(
self.context.auth_server_url, self.context.server_url
)
# Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers)
for url in asm_discovery_urls: # pragma: no cover
oauth_metadata_request = create_oauth_metadata_request(url)
oauth_metadata_response = yield oauth_metadata_request
ok, asm = await handle_auth_metadata_response(oauth_metadata_response)
if not ok:
break
if ok and asm:
self.context.oauth_metadata = asm
break
else:
logger.debug(f"OAuth metadata discovery failed: {url}")
# Step 3: Apply scope selection strategy
self.context.client_metadata.scope = get_client_metadata_scopes(
extract_scope_from_www_auth(response),
self.context.protected_resource_metadata,
self.context.oauth_metadata,
)
# Step 4: Register client or use URL-based client ID (CIMD)
if not self.context.client_info:
if should_use_client_metadata_url(
self.context.oauth_metadata, self.context.client_metadata_url
):
# Use URL-based client ID (CIMD)
logger.debug(f"Using URL-based client ID (CIMD): {self.context.client_metadata_url}")
client_information = create_client_info_from_metadata_url(
self.context.client_metadata_url, # type: ignore[arg-type]
redirect_uris=self.context.client_metadata.redirect_uris,
)
self.context.client_info = client_information
await self.context.storage.set_client_info(client_information)
else:
# Fallback to Dynamic Client Registration
registration_request = create_client_registration_request(
self.context.oauth_metadata,
self.context.client_metadata,
self.context.get_authorization_base_url(self.context.server_url),
)
registration_response = yield registration_request
client_information = await handle_registration_response(registration_response)
self.context.client_info = client_information
await self.context.storage.set_client_info(client_information)
# Step 5: Perform authorization and complete token exchange
token_response = yield await self._perform_authorization()
await self._handle_token_response(token_response)
except Exception: # pragma: no cover
logger.exception("OAuth flow error")
raise
# Retry with new tokens
self._add_auth_header(request)
yield request
elif response.status_code == 403:
# Step 1: Extract error field from WWW-Authenticate header
error = extract_field_from_www_auth(response, "error")
# Step 2: Check if we need to step-up authorization
if error == "insufficient_scope": # pragma: no branch
try:
# Step 2a: Update the required scopes
self.context.client_metadata.scope = get_client_metadata_scopes(
extract_scope_from_www_auth(response), self.context.protected_resource_metadata
)
# Step 2b: Perform (re-)authorization and token exchange
token_response = yield await self._perform_authorization()
await self._handle_token_response(token_response)
except Exception: # pragma: no cover
logger.exception("OAuth flow error")
raise
# Retry with new tokens
self._add_auth_header(request)
yield request
@@ -0,0 +1,336 @@
import logging
import re
from urllib.parse import urljoin, urlparse
from httpx import Request, Response
from pydantic import AnyUrl, ValidationError
from mcp.client.auth import OAuthRegistrationError, OAuthTokenError
from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
from mcp.shared.auth import (
OAuthClientInformationFull,
OAuthClientMetadata,
OAuthMetadata,
OAuthToken,
ProtectedResourceMetadata,
)
from mcp.types import LATEST_PROTOCOL_VERSION
logger = logging.getLogger(__name__)
def extract_field_from_www_auth(response: Response, field_name: str) -> str | None:
"""
Extract field from WWW-Authenticate header.
Returns:
Field value if found in WWW-Authenticate header, None otherwise
"""
www_auth_header = response.headers.get("WWW-Authenticate")
if not www_auth_header:
return None
# Pattern matches: field_name="value" or field_name=value (unquoted)
pattern = rf'{field_name}=(?:"([^"]+)"|([^\s,]+))'
match = re.search(pattern, www_auth_header)
if match:
# Return quoted value if present, otherwise unquoted value
return match.group(1) or match.group(2)
return None
def extract_scope_from_www_auth(response: Response) -> str | None:
"""
Extract scope parameter from WWW-Authenticate header as per RFC6750.
Returns:
Scope string if found in WWW-Authenticate header, None otherwise
"""
return extract_field_from_www_auth(response, "scope")
def extract_resource_metadata_from_www_auth(response: Response) -> str | None:
"""
Extract protected resource metadata URL from WWW-Authenticate header as per RFC9728.
Returns:
Resource metadata URL if found in WWW-Authenticate header, None otherwise
"""
if not response or response.status_code != 401:
return None # pragma: no cover
return extract_field_from_www_auth(response, "resource_metadata")
def build_protected_resource_metadata_discovery_urls(www_auth_url: str | None, server_url: str) -> list[str]:
"""
Build ordered list of URLs to try for protected resource metadata discovery.
Per SEP-985, the client MUST:
1. Try resource_metadata from WWW-Authenticate header (if present)
2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path}
3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource
Args:
www_auth_url: optional resource_metadata url extracted from the WWW-Authenticate header
server_url: server url
Returns:
Ordered list of URLs to try for discovery
"""
urls: list[str] = []
# Priority 1: WWW-Authenticate header with resource_metadata parameter
if www_auth_url:
urls.append(www_auth_url)
# Priority 2-3: Well-known URIs (RFC 9728)
parsed = urlparse(server_url)
base_url = f"{parsed.scheme}://{parsed.netloc}"
# Priority 2: Path-based well-known URI (if server has a path component)
if parsed.path and parsed.path != "/":
path_based_url = urljoin(base_url, f"/.well-known/oauth-protected-resource{parsed.path}")
urls.append(path_based_url)
# Priority 3: Root-based well-known URI
root_based_url = urljoin(base_url, "/.well-known/oauth-protected-resource")
urls.append(root_based_url)
return urls
def get_client_metadata_scopes(
www_authenticate_scope: str | None,
protected_resource_metadata: ProtectedResourceMetadata | None,
authorization_server_metadata: OAuthMetadata | None = None,
) -> str | None:
"""Select scopes as outlined in the 'Scope Selection Strategy' in the MCP spec."""
# Per MCP spec, scope selection priority order:
# 1. Use scope from WWW-Authenticate header (if provided)
# 2. Use all scopes from PRM scopes_supported (if available)
# 3. Omit scope parameter if neither is available
if www_authenticate_scope is not None:
# Priority 1: WWW-Authenticate header scope
return www_authenticate_scope
elif protected_resource_metadata is not None and protected_resource_metadata.scopes_supported is not None:
# Priority 2: PRM scopes_supported
return " ".join(protected_resource_metadata.scopes_supported)
elif authorization_server_metadata is not None and authorization_server_metadata.scopes_supported is not None:
return " ".join(authorization_server_metadata.scopes_supported) # pragma: no cover
else:
# Priority 3: Omit scope parameter
return None
def build_oauth_authorization_server_metadata_discovery_urls(auth_server_url: str | None, server_url: str) -> list[str]:
"""
Generate ordered list of (url, type) tuples for discovery attempts.
Args:
auth_server_url: URL for the OAuth Authorization Metadata URL if found, otherwise None
server_url: URL for the MCP server, used as a fallback if auth_server_url is None
"""
if not auth_server_url:
# Legacy path using the 2025-03-26 spec:
# link: https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization
parsed = urlparse(server_url)
return [f"{parsed.scheme}://{parsed.netloc}/.well-known/oauth-authorization-server"]
urls: list[str] = []
parsed = urlparse(auth_server_url)
base_url = f"{parsed.scheme}://{parsed.netloc}"
# RFC 8414: Path-aware OAuth discovery
if parsed.path and parsed.path != "/":
oauth_path = f"/.well-known/oauth-authorization-server{parsed.path.rstrip('/')}"
urls.append(urljoin(base_url, oauth_path))
# RFC 8414 section 5: Path-aware OIDC discovery
# See https://www.rfc-editor.org/rfc/rfc8414.html#section-5
oidc_path = f"/.well-known/openid-configuration{parsed.path.rstrip('/')}"
urls.append(urljoin(base_url, oidc_path))
# https://openid.net/specs/openid-connect-discovery-1_0.html
oidc_path = f"{parsed.path.rstrip('/')}/.well-known/openid-configuration"
urls.append(urljoin(base_url, oidc_path))
return urls
# OAuth root
urls.append(urljoin(base_url, "/.well-known/oauth-authorization-server"))
# OIDC 1.0 fallback (appends to full URL per OIDC spec)
# https://openid.net/specs/openid-connect-discovery-1_0.html
urls.append(urljoin(base_url, "/.well-known/openid-configuration"))
return urls
async def handle_protected_resource_response(
response: Response,
) -> ProtectedResourceMetadata | None:
"""
Handle protected resource metadata discovery response.
Per SEP-985, supports fallback when discovery fails at one URL.
Returns:
True if metadata was successfully discovered, False if we should try next URL
"""
if response.status_code == 200:
try:
content = await response.aread()
metadata = ProtectedResourceMetadata.model_validate_json(content)
return metadata
except ValidationError: # pragma: no cover
# Invalid metadata - try next URL
return None
else:
# Not found - try next URL in fallback chain
return None
async def handle_auth_metadata_response(response: Response) -> tuple[bool, OAuthMetadata | None]:
if response.status_code == 200:
try:
content = await response.aread()
asm = OAuthMetadata.model_validate_json(content)
return True, asm
except ValidationError: # pragma: no cover
return True, None
elif response.status_code < 400 or response.status_code >= 500:
return False, None # Non-4XX error, stop trying
return True, None
def create_oauth_metadata_request(url: str) -> Request:
return Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
def create_client_registration_request(
auth_server_metadata: OAuthMetadata | None, client_metadata: OAuthClientMetadata, auth_base_url: str
) -> Request:
"""Build registration request or skip if already registered."""
if auth_server_metadata and auth_server_metadata.registration_endpoint:
registration_url = str(auth_server_metadata.registration_endpoint)
else:
registration_url = urljoin(auth_base_url, "/register")
registration_data = client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True)
return Request("POST", registration_url, json=registration_data, headers={"Content-Type": "application/json"})
async def handle_registration_response(response: Response) -> OAuthClientInformationFull:
"""Handle registration response."""
if response.status_code not in (200, 201):
await response.aread()
raise OAuthRegistrationError(f"Registration failed: {response.status_code} {response.text}")
try:
content = await response.aread()
client_info = OAuthClientInformationFull.model_validate_json(content)
return client_info
# self.context.client_info = client_info
# await self.context.storage.set_client_info(client_info)
except ValidationError as e: # pragma: no cover
raise OAuthRegistrationError(f"Invalid registration response: {e}")
def is_valid_client_metadata_url(url: str | None) -> bool:
"""Validate that a URL is suitable for use as a client_id (CIMD).
The URL must be HTTPS with a non-root pathname.
Args:
url: The URL to validate
Returns:
True if the URL is a valid HTTPS URL with a non-root pathname
"""
if not url:
return False
try:
parsed = urlparse(url)
return parsed.scheme == "https" and parsed.path not in ("", "/")
except Exception:
return False
def should_use_client_metadata_url(
oauth_metadata: OAuthMetadata | None,
client_metadata_url: str | None,
) -> bool:
"""Determine if URL-based client ID (CIMD) should be used instead of DCR.
URL-based client IDs should be used when:
1. The server advertises client_id_metadata_document_supported=true
2. The client has a valid client_metadata_url configured
Args:
oauth_metadata: OAuth authorization server metadata
client_metadata_url: URL-based client ID (already validated)
Returns:
True if CIMD should be used, False if DCR should be used
"""
if not client_metadata_url:
return False
if not oauth_metadata:
return False
return oauth_metadata.client_id_metadata_document_supported is True
def create_client_info_from_metadata_url(
client_metadata_url: str, redirect_uris: list[AnyUrl] | None = None
) -> OAuthClientInformationFull:
"""Create client information using a URL-based client ID (CIMD).
When using URL-based client IDs, the URL itself becomes the client_id
and no client_secret is used (token_endpoint_auth_method="none").
Args:
client_metadata_url: The URL to use as the client_id
redirect_uris: The redirect URIs from the client metadata (passed through for
compatibility with OAuthClientInformationFull which inherits from OAuthClientMetadata)
Returns:
OAuthClientInformationFull with the URL as client_id
"""
return OAuthClientInformationFull(
client_id=client_metadata_url,
token_endpoint_auth_method="none",
redirect_uris=redirect_uris,
)
async def handle_token_response_scopes(
response: Response,
) -> OAuthToken:
"""Parse and validate token response with optional scope validation.
Parses token response JSON. Callers should check response.status_code before calling.
Args:
response: HTTP response from token endpoint (status already checked by caller)
Returns:
Validated OAuthToken model
Raises:
OAuthTokenError: If response JSON is invalid
"""
try:
content = await response.aread()
token_response = OAuthToken.model_validate_json(content)
return token_response
except ValidationError as e: # pragma: no cover
raise OAuthTokenError(f"Invalid token response: {e}")
@@ -0,0 +1,9 @@
"""
Experimental client features.
WARNING: These APIs are experimental and may change without notice.
"""
from mcp.client.experimental.tasks import ExperimentalClientFeatures
__all__ = ["ExperimentalClientFeatures"]
@@ -0,0 +1,290 @@
"""
Experimental task handler protocols for server -> client requests.
This module provides Protocol types and default handlers for when servers
send task-related requests to clients (the reverse of normal client -> server flow).
WARNING: These APIs are experimental and may change without notice.
Use cases:
- Server sends task-augmented sampling/elicitation request to client
- Client creates a local task, spawns background work, returns CreateTaskResult
- Server polls client's task status via tasks/get, tasks/result, etc.
"""
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol
from pydantic import TypeAdapter
import mcp.types as types
from mcp.shared.context import RequestContext
from mcp.shared.session import RequestResponder
if TYPE_CHECKING:
from mcp.client.session import ClientSession
class GetTaskHandlerFnT(Protocol):
"""Handler for tasks/get requests from server.
WARNING: This is experimental and may change without notice.
"""
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.GetTaskRequestParams,
) -> types.GetTaskResult | types.ErrorData: ... # pragma: no branch
class GetTaskResultHandlerFnT(Protocol):
"""Handler for tasks/result requests from server.
WARNING: This is experimental and may change without notice.
"""
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.GetTaskPayloadRequestParams,
) -> types.GetTaskPayloadResult | types.ErrorData: ... # pragma: no branch
class ListTasksHandlerFnT(Protocol):
"""Handler for tasks/list requests from server.
WARNING: This is experimental and may change without notice.
"""
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData: ... # pragma: no branch
class CancelTaskHandlerFnT(Protocol):
"""Handler for tasks/cancel requests from server.
WARNING: This is experimental and may change without notice.
"""
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData: ... # pragma: no branch
class TaskAugmentedSamplingFnT(Protocol):
"""Handler for task-augmented sampling/createMessage requests from server.
When server sends a CreateMessageRequest with task field, this callback
is invoked. The callback should create a task, spawn background work,
and return CreateTaskResult immediately.
WARNING: This is experimental and may change without notice.
"""
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.CreateMessageRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData: ... # pragma: no branch
class TaskAugmentedElicitationFnT(Protocol):
"""Handler for task-augmented elicitation/create requests from server.
When server sends an ElicitRequest with task field, this callback
is invoked. The callback should create a task, spawn background work,
and return CreateTaskResult immediately.
WARNING: This is experimental and may change without notice.
"""
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.ElicitRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData: ... # pragma: no branch
async def default_get_task_handler(
context: RequestContext["ClientSession", Any],
params: types.GetTaskRequestParams,
) -> types.GetTaskResult | types.ErrorData:
return types.ErrorData(
code=types.METHOD_NOT_FOUND,
message="tasks/get not supported",
)
async def default_get_task_result_handler(
context: RequestContext["ClientSession", Any],
params: types.GetTaskPayloadRequestParams,
) -> types.GetTaskPayloadResult | types.ErrorData:
return types.ErrorData(
code=types.METHOD_NOT_FOUND,
message="tasks/result not supported",
)
async def default_list_tasks_handler(
context: RequestContext["ClientSession", Any],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData:
return types.ErrorData(
code=types.METHOD_NOT_FOUND,
message="tasks/list not supported",
)
async def default_cancel_task_handler(
context: RequestContext["ClientSession", Any],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData:
return types.ErrorData(
code=types.METHOD_NOT_FOUND,
message="tasks/cancel not supported",
)
async def default_task_augmented_sampling(
context: RequestContext["ClientSession", Any],
params: types.CreateMessageRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData:
return types.ErrorData(
code=types.INVALID_REQUEST,
message="Task-augmented sampling not supported",
)
async def default_task_augmented_elicitation(
context: RequestContext["ClientSession", Any],
params: types.ElicitRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData:
return types.ErrorData(
code=types.INVALID_REQUEST,
message="Task-augmented elicitation not supported",
)
@dataclass
class ExperimentalTaskHandlers:
"""Container for experimental task handlers.
Groups all task-related handlers that handle server -> client requests.
This includes both pure task requests (get, list, cancel, result) and
task-augmented request handlers (sampling, elicitation with task field).
WARNING: These APIs are experimental and may change without notice.
Example:
handlers = ExperimentalTaskHandlers(
get_task=my_get_task_handler,
list_tasks=my_list_tasks_handler,
)
session = ClientSession(..., experimental_task_handlers=handlers)
"""
# Pure task request handlers
get_task: GetTaskHandlerFnT = field(default=default_get_task_handler)
get_task_result: GetTaskResultHandlerFnT = field(default=default_get_task_result_handler)
list_tasks: ListTasksHandlerFnT = field(default=default_list_tasks_handler)
cancel_task: CancelTaskHandlerFnT = field(default=default_cancel_task_handler)
# Task-augmented request handlers
augmented_sampling: TaskAugmentedSamplingFnT = field(default=default_task_augmented_sampling)
augmented_elicitation: TaskAugmentedElicitationFnT = field(default=default_task_augmented_elicitation)
def build_capability(self) -> types.ClientTasksCapability | None:
"""Build ClientTasksCapability from the configured handlers.
Returns a capability object that reflects which handlers are configured
(i.e., not using the default "not supported" handlers).
Returns:
ClientTasksCapability if any handlers are provided, None otherwise
"""
has_list = self.list_tasks is not default_list_tasks_handler
has_cancel = self.cancel_task is not default_cancel_task_handler
has_sampling = self.augmented_sampling is not default_task_augmented_sampling
has_elicitation = self.augmented_elicitation is not default_task_augmented_elicitation
# If no handlers are provided, return None
if not any([has_list, has_cancel, has_sampling, has_elicitation]):
return None
# Build requests capability if any request handlers are provided
requests_capability: types.ClientTasksRequestsCapability | None = None
if has_sampling or has_elicitation:
requests_capability = types.ClientTasksRequestsCapability(
sampling=types.TasksSamplingCapability(createMessage=types.TasksCreateMessageCapability())
if has_sampling
else None,
elicitation=types.TasksElicitationCapability(create=types.TasksCreateElicitationCapability())
if has_elicitation
else None,
)
return types.ClientTasksCapability(
list=types.TasksListCapability() if has_list else None,
cancel=types.TasksCancelCapability() if has_cancel else None,
requests=requests_capability,
)
@staticmethod
def handles_request(request: types.ServerRequest) -> bool:
"""Check if this handler handles the given request type."""
return isinstance(
request.root,
types.GetTaskRequest | types.GetTaskPayloadRequest | types.ListTasksRequest | types.CancelTaskRequest,
)
async def handle_request(
self,
ctx: RequestContext["ClientSession", Any],
responder: RequestResponder[types.ServerRequest, types.ClientResult],
) -> None:
"""Handle a task-related request from the server.
Call handles_request() first to check if this handler can handle the request.
"""
client_response_type: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(
types.ClientResult | types.ErrorData
)
match responder.request.root:
case types.GetTaskRequest(params=params):
response = await self.get_task(ctx, params)
client_response = client_response_type.validate_python(response)
await responder.respond(client_response)
case types.GetTaskPayloadRequest(params=params):
response = await self.get_task_result(ctx, params)
client_response = client_response_type.validate_python(response)
await responder.respond(client_response)
case types.ListTasksRequest(params=params):
response = await self.list_tasks(ctx, params)
client_response = client_response_type.validate_python(response)
await responder.respond(client_response)
case types.CancelTaskRequest(params=params):
response = await self.cancel_task(ctx, params)
client_response = client_response_type.validate_python(response)
await responder.respond(client_response)
case _: # pragma: no cover
raise ValueError(f"Unhandled request type: {type(responder.request.root)}")
# Backwards compatibility aliases
default_task_augmented_sampling_callback = default_task_augmented_sampling
default_task_augmented_elicitation_callback = default_task_augmented_elicitation
@@ -0,0 +1,224 @@
"""
Experimental client-side task support.
This module provides client methods for interacting with MCP tasks.
WARNING: These APIs are experimental and may change without notice.
Example:
# Call a tool as a task
result = await session.experimental.call_tool_as_task("tool_name", {"arg": "value"})
task_id = result.task.taskId
# Get task status
status = await session.experimental.get_task(task_id)
# Get task result when complete
if status.status == "completed":
result = await session.experimental.get_task_result(task_id, CallToolResult)
# List all tasks
tasks = await session.experimental.list_tasks()
# Cancel a task
await session.experimental.cancel_task(task_id)
"""
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Any, TypeVar
import mcp.types as types
from mcp.shared.experimental.tasks.polling import poll_until_terminal
if TYPE_CHECKING:
from mcp.client.session import ClientSession
ResultT = TypeVar("ResultT", bound=types.Result)
class ExperimentalClientFeatures:
"""
Experimental client features for tasks and other experimental APIs.
WARNING: These APIs are experimental and may change without notice.
Access via session.experimental:
status = await session.experimental.get_task(task_id)
"""
def __init__(self, session: "ClientSession") -> None:
self._session = session
async def call_tool_as_task(
self,
name: str,
arguments: dict[str, Any] | None = None,
*,
ttl: int = 60000,
meta: dict[str, Any] | None = None,
) -> types.CreateTaskResult:
"""Call a tool as a task, returning a CreateTaskResult for polling.
This is a convenience method for calling tools that support task execution.
The server will return a task reference instead of the immediate result,
which can then be polled via `get_task()` and retrieved via `get_task_result()`.
Args:
name: The tool name
arguments: Tool arguments
ttl: Task time-to-live in milliseconds (default: 60000 = 1 minute)
meta: Optional metadata to include in the request
Returns:
CreateTaskResult containing the task reference
Example:
# Create task
result = await session.experimental.call_tool_as_task(
"long_running_tool", {"input": "data"}
)
task_id = result.task.taskId
# Poll for completion
while True:
status = await session.experimental.get_task(task_id)
if status.status == "completed":
break
await asyncio.sleep(0.5)
# Get result
final = await session.experimental.get_task_result(task_id, CallToolResult)
"""
_meta: types.RequestParams.Meta | None = None
if meta is not None:
_meta = types.RequestParams.Meta(**meta)
return await self._session.send_request(
types.ClientRequest(
types.CallToolRequest(
params=types.CallToolRequestParams(
name=name,
arguments=arguments,
task=types.TaskMetadata(ttl=ttl),
_meta=_meta,
),
)
),
types.CreateTaskResult,
)
async def get_task(self, task_id: str) -> types.GetTaskResult:
"""
Get the current status of a task.
Args:
task_id: The task identifier
Returns:
GetTaskResult containing the task status and metadata
"""
return await self._session.send_request(
types.ClientRequest(
types.GetTaskRequest(
params=types.GetTaskRequestParams(taskId=task_id),
)
),
types.GetTaskResult,
)
async def get_task_result(
self,
task_id: str,
result_type: type[ResultT],
) -> ResultT:
"""
Get the result of a completed task.
The result type depends on the original request type:
- tools/call tasks return CallToolResult
- Other request types return their corresponding result type
Args:
task_id: The task identifier
result_type: The expected result type (e.g., CallToolResult)
Returns:
The task result, validated against result_type
"""
return await self._session.send_request(
types.ClientRequest(
types.GetTaskPayloadRequest(
params=types.GetTaskPayloadRequestParams(taskId=task_id),
)
),
result_type,
)
async def list_tasks(
self,
cursor: str | None = None,
) -> types.ListTasksResult:
"""
List all tasks.
Args:
cursor: Optional pagination cursor
Returns:
ListTasksResult containing tasks and optional next cursor
"""
params = types.PaginatedRequestParams(cursor=cursor) if cursor else None
return await self._session.send_request(
types.ClientRequest(
types.ListTasksRequest(params=params),
),
types.ListTasksResult,
)
async def cancel_task(self, task_id: str) -> types.CancelTaskResult:
"""
Cancel a running task.
Args:
task_id: The task identifier
Returns:
CancelTaskResult with the updated task state
"""
return await self._session.send_request(
types.ClientRequest(
types.CancelTaskRequest(
params=types.CancelTaskRequestParams(taskId=task_id),
)
),
types.CancelTaskResult,
)
async def poll_task(self, task_id: str) -> AsyncIterator[types.GetTaskResult]:
"""
Poll a task until it reaches a terminal status.
Yields GetTaskResult for each poll, allowing the caller to react to
status changes (e.g., handle input_required). Exits when task reaches
a terminal status (completed, failed, cancelled).
Respects the pollInterval hint from the server.
Args:
task_id: The task identifier
Yields:
GetTaskResult for each poll
Example:
async for status in session.experimental.poll_task(task_id):
print(f"Status: {status.status}")
if status.status == "input_required":
# Handle elicitation request via tasks/result
pass
# Task is now terminal, get the result
result = await session.experimental.get_task_result(task_id, CallToolResult)
"""
async for status in poll_until_terminal(self.get_task, task_id):
yield status
@@ -0,0 +1,615 @@
import logging
from datetime import timedelta
from typing import Any, Protocol, overload
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import AnyUrl, TypeAdapter
from typing_extensions import deprecated
import mcp.types as types
from mcp.client.experimental import ExperimentalClientFeatures
from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
from mcp.shared.context import RequestContext
from mcp.shared.message import SessionMessage
from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
logger = logging.getLogger("client")
class SamplingFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: ... # pragma: no branch
class ElicitationFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
params: types.ElicitRequestParams,
) -> types.ElicitResult | types.ErrorData: ... # pragma: no branch
class ListRootsFnT(Protocol):
async def __call__(
self, context: RequestContext["ClientSession", Any]
) -> types.ListRootsResult | types.ErrorData: ... # pragma: no branch
class LoggingFnT(Protocol):
async def __call__(
self,
params: types.LoggingMessageNotificationParams,
) -> None: ... # pragma: no branch
class MessageHandlerFnT(Protocol):
async def __call__(
self,
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
) -> None: ... # pragma: no branch
async def _default_message_handler(
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
) -> None:
await anyio.lowlevel.checkpoint()
async def _default_sampling_callback(
context: RequestContext["ClientSession", Any],
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData:
return types.ErrorData(
code=types.INVALID_REQUEST,
message="Sampling not supported",
)
async def _default_elicitation_callback(
context: RequestContext["ClientSession", Any],
params: types.ElicitRequestParams,
) -> types.ElicitResult | types.ErrorData:
return types.ErrorData( # pragma: no cover
code=types.INVALID_REQUEST,
message="Elicitation not supported",
)
async def _default_list_roots_callback(
context: RequestContext["ClientSession", Any],
) -> types.ListRootsResult | types.ErrorData:
return types.ErrorData(
code=types.INVALID_REQUEST,
message="List roots not supported",
)
async def _default_logging_callback(
params: types.LoggingMessageNotificationParams,
) -> None:
pass
ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData)
class ClientSession(
BaseSession[
types.ClientRequest,
types.ClientNotification,
types.ClientResult,
types.ServerRequest,
types.ServerNotification,
]
):
def __init__(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_timeout_seconds: timedelta | None = None,
sampling_callback: SamplingFnT | None = None,
elicitation_callback: ElicitationFnT | None = None,
list_roots_callback: ListRootsFnT | None = None,
logging_callback: LoggingFnT | None = None,
message_handler: MessageHandlerFnT | None = None,
client_info: types.Implementation | None = None,
*,
sampling_capabilities: types.SamplingCapability | None = None,
experimental_task_handlers: ExperimentalTaskHandlers | None = None,
) -> None:
super().__init__(
read_stream,
write_stream,
types.ServerRequest,
types.ServerNotification,
read_timeout_seconds=read_timeout_seconds,
)
self._client_info = client_info or DEFAULT_CLIENT_INFO
self._sampling_callback = sampling_callback or _default_sampling_callback
self._sampling_capabilities = sampling_capabilities
self._elicitation_callback = elicitation_callback or _default_elicitation_callback
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
self._logging_callback = logging_callback or _default_logging_callback
self._message_handler = message_handler or _default_message_handler
self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
self._server_capabilities: types.ServerCapabilities | None = None
self._experimental_features: ExperimentalClientFeatures | None = None
# Experimental: Task handlers (use defaults if not provided)
self._task_handlers = experimental_task_handlers or ExperimentalTaskHandlers()
async def initialize(self) -> types.InitializeResult:
sampling = (
(self._sampling_capabilities or types.SamplingCapability())
if self._sampling_callback is not _default_sampling_callback
else None
)
elicitation = (
types.ElicitationCapability(
form=types.FormElicitationCapability(),
url=types.UrlElicitationCapability(),
)
if self._elicitation_callback is not _default_elicitation_callback
else None
)
roots = (
# TODO: Should this be based on whether we
# _will_ send notifications, or only whether
# they're supported?
types.RootsCapability(listChanged=True)
if self._list_roots_callback is not _default_list_roots_callback
else None
)
result = await self.send_request(
types.ClientRequest(
types.InitializeRequest(
params=types.InitializeRequestParams(
protocolVersion=types.LATEST_PROTOCOL_VERSION,
capabilities=types.ClientCapabilities(
sampling=sampling,
elicitation=elicitation,
experimental=None,
roots=roots,
tasks=self._task_handlers.build_capability(),
),
clientInfo=self._client_info,
),
)
),
types.InitializeResult,
)
if result.protocolVersion not in SUPPORTED_PROTOCOL_VERSIONS:
raise RuntimeError(f"Unsupported protocol version from the server: {result.protocolVersion}")
self._server_capabilities = result.capabilities
await self.send_notification(types.ClientNotification(types.InitializedNotification()))
return result
def get_server_capabilities(self) -> types.ServerCapabilities | None:
"""Return the server capabilities received during initialization.
Returns None if the session has not been initialized yet.
"""
return self._server_capabilities
@property
def experimental(self) -> ExperimentalClientFeatures:
"""Experimental APIs for tasks and other features.
WARNING: These APIs are experimental and may change without notice.
Example:
status = await session.experimental.get_task(task_id)
result = await session.experimental.get_task_result(task_id, CallToolResult)
"""
if self._experimental_features is None:
self._experimental_features = ExperimentalClientFeatures(self)
return self._experimental_features
async def send_ping(self) -> types.EmptyResult:
"""Send a ping request."""
return await self.send_request(
types.ClientRequest(types.PingRequest()),
types.EmptyResult,
)
async def send_progress_notification(
self,
progress_token: str | int,
progress: float,
total: float | None = None,
message: str | None = None,
) -> None:
"""Send a progress notification."""
await self.send_notification(
types.ClientNotification(
types.ProgressNotification(
params=types.ProgressNotificationParams(
progressToken=progress_token,
progress=progress,
total=total,
message=message,
),
),
)
)
async def set_logging_level(self, level: types.LoggingLevel) -> types.EmptyResult:
"""Send a logging/setLevel request."""
return await self.send_request( # pragma: no cover
types.ClientRequest(
types.SetLevelRequest(
params=types.SetLevelRequestParams(level=level),
)
),
types.EmptyResult,
)
@overload
@deprecated("Use list_resources(params=PaginatedRequestParams(...)) instead")
async def list_resources(self, cursor: str | None) -> types.ListResourcesResult: ...
@overload
async def list_resources(self, *, params: types.PaginatedRequestParams | None) -> types.ListResourcesResult: ...
@overload
async def list_resources(self) -> types.ListResourcesResult: ...
async def list_resources(
self,
cursor: str | None = None,
*,
params: types.PaginatedRequestParams | None = None,
) -> types.ListResourcesResult:
"""Send a resources/list request.
Args:
cursor: Simple cursor string for pagination (deprecated, use params instead)
params: Full pagination parameters including cursor and any future fields
"""
if params is not None and cursor is not None:
raise ValueError("Cannot specify both cursor and params")
if params is not None:
request_params = params
elif cursor is not None:
request_params = types.PaginatedRequestParams(cursor=cursor)
else:
request_params = None
return await self.send_request(
types.ClientRequest(types.ListResourcesRequest(params=request_params)),
types.ListResourcesResult,
)
@overload
@deprecated("Use list_resource_templates(params=PaginatedRequestParams(...)) instead")
async def list_resource_templates(self, cursor: str | None) -> types.ListResourceTemplatesResult: ...
@overload
async def list_resource_templates(
self, *, params: types.PaginatedRequestParams | None
) -> types.ListResourceTemplatesResult: ...
@overload
async def list_resource_templates(self) -> types.ListResourceTemplatesResult: ...
async def list_resource_templates(
self,
cursor: str | None = None,
*,
params: types.PaginatedRequestParams | None = None,
) -> types.ListResourceTemplatesResult:
"""Send a resources/templates/list request.
Args:
cursor: Simple cursor string for pagination (deprecated, use params instead)
params: Full pagination parameters including cursor and any future fields
"""
if params is not None and cursor is not None:
raise ValueError("Cannot specify both cursor and params")
if params is not None:
request_params = params
elif cursor is not None:
request_params = types.PaginatedRequestParams(cursor=cursor)
else:
request_params = None
return await self.send_request(
types.ClientRequest(types.ListResourceTemplatesRequest(params=request_params)),
types.ListResourceTemplatesResult,
)
async def read_resource(self, uri: AnyUrl) -> types.ReadResourceResult:
"""Send a resources/read request."""
return await self.send_request(
types.ClientRequest(
types.ReadResourceRequest(
params=types.ReadResourceRequestParams(uri=uri),
)
),
types.ReadResourceResult,
)
async def subscribe_resource(self, uri: AnyUrl) -> types.EmptyResult:
"""Send a resources/subscribe request."""
return await self.send_request( # pragma: no cover
types.ClientRequest(
types.SubscribeRequest(
params=types.SubscribeRequestParams(uri=uri),
)
),
types.EmptyResult,
)
async def unsubscribe_resource(self, uri: AnyUrl) -> types.EmptyResult:
"""Send a resources/unsubscribe request."""
return await self.send_request( # pragma: no cover
types.ClientRequest(
types.UnsubscribeRequest(
params=types.UnsubscribeRequestParams(uri=uri),
)
),
types.EmptyResult,
)
async def call_tool(
self,
name: str,
arguments: dict[str, Any] | None = None,
read_timeout_seconds: timedelta | None = None,
progress_callback: ProgressFnT | None = None,
*,
meta: dict[str, Any] | None = None,
) -> types.CallToolResult:
"""Send a tools/call request with optional progress callback support."""
_meta: types.RequestParams.Meta | None = None
if meta is not None:
_meta = types.RequestParams.Meta(**meta)
result = await self.send_request(
types.ClientRequest(
types.CallToolRequest(
params=types.CallToolRequestParams(name=name, arguments=arguments, _meta=_meta),
)
),
types.CallToolResult,
request_read_timeout_seconds=read_timeout_seconds,
progress_callback=progress_callback,
)
if not result.isError:
await self._validate_tool_result(name, result)
return result
async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
"""Validate the structured content of a tool result against its output schema."""
if name not in self._tool_output_schemas:
# refresh output schema cache
await self.list_tools()
output_schema = None
if name in self._tool_output_schemas:
output_schema = self._tool_output_schemas.get(name)
else:
logger.warning(f"Tool {name} not listed by server, cannot validate any structured content")
if output_schema is not None:
from jsonschema import SchemaError, ValidationError, validate
if result.structuredContent is None:
raise RuntimeError(
f"Tool {name} has an output schema but did not return structured content"
) # pragma: no cover
try:
validate(result.structuredContent, output_schema)
except ValidationError as e:
raise RuntimeError(f"Invalid structured content returned by tool {name}: {e}") # pragma: no cover
except SchemaError as e: # pragma: no cover
raise RuntimeError(f"Invalid schema for tool {name}: {e}") # pragma: no cover
@overload
@deprecated("Use list_prompts(params=PaginatedRequestParams(...)) instead")
async def list_prompts(self, cursor: str | None) -> types.ListPromptsResult: ...
@overload
async def list_prompts(self, *, params: types.PaginatedRequestParams | None) -> types.ListPromptsResult: ...
@overload
async def list_prompts(self) -> types.ListPromptsResult: ...
async def list_prompts(
self,
cursor: str | None = None,
*,
params: types.PaginatedRequestParams | None = None,
) -> types.ListPromptsResult:
"""Send a prompts/list request.
Args:
cursor: Simple cursor string for pagination (deprecated, use params instead)
params: Full pagination parameters including cursor and any future fields
"""
if params is not None and cursor is not None:
raise ValueError("Cannot specify both cursor and params")
if params is not None:
request_params = params
elif cursor is not None:
request_params = types.PaginatedRequestParams(cursor=cursor)
else:
request_params = None
return await self.send_request(
types.ClientRequest(types.ListPromptsRequest(params=request_params)),
types.ListPromptsResult,
)
async def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> types.GetPromptResult:
"""Send a prompts/get request."""
return await self.send_request(
types.ClientRequest(
types.GetPromptRequest(
params=types.GetPromptRequestParams(name=name, arguments=arguments),
)
),
types.GetPromptResult,
)
async def complete(
self,
ref: types.ResourceTemplateReference | types.PromptReference,
argument: dict[str, str],
context_arguments: dict[str, str] | None = None,
) -> types.CompleteResult:
"""Send a completion/complete request."""
context = None
if context_arguments is not None:
context = types.CompletionContext(arguments=context_arguments)
return await self.send_request(
types.ClientRequest(
types.CompleteRequest(
params=types.CompleteRequestParams(
ref=ref,
argument=types.CompletionArgument(**argument),
context=context,
),
)
),
types.CompleteResult,
)
@overload
@deprecated("Use list_tools(params=PaginatedRequestParams(...)) instead")
async def list_tools(self, cursor: str | None) -> types.ListToolsResult: ...
@overload
async def list_tools(self, *, params: types.PaginatedRequestParams | None) -> types.ListToolsResult: ...
@overload
async def list_tools(self) -> types.ListToolsResult: ...
async def list_tools(
self,
cursor: str | None = None,
*,
params: types.PaginatedRequestParams | None = None,
) -> types.ListToolsResult:
"""Send a tools/list request.
Args:
cursor: Simple cursor string for pagination (deprecated, use params instead)
params: Full pagination parameters including cursor and any future fields
"""
if params is not None and cursor is not None:
raise ValueError("Cannot specify both cursor and params")
if params is not None:
request_params = params
elif cursor is not None:
request_params = types.PaginatedRequestParams(cursor=cursor)
else:
request_params = None
result = await self.send_request(
types.ClientRequest(types.ListToolsRequest(params=request_params)),
types.ListToolsResult,
)
# Cache tool output schemas for future validation
# Note: don't clear the cache, as we may be using a cursor
for tool in result.tools:
self._tool_output_schemas[tool.name] = tool.outputSchema
return result
async def send_roots_list_changed(self) -> None: # pragma: no cover
"""Send a roots/list_changed notification."""
await self.send_notification(types.ClientNotification(types.RootsListChangedNotification()))
async def _received_request(self, responder: RequestResponder[types.ServerRequest, types.ClientResult]) -> None:
ctx = RequestContext[ClientSession, Any](
request_id=responder.request_id,
meta=responder.request_meta,
session=self,
lifespan_context=None,
)
# Delegate to experimental task handler if applicable
if self._task_handlers.handles_request(responder.request):
with responder:
await self._task_handlers.handle_request(ctx, responder)
return None
# Core request handling
match responder.request.root:
case types.CreateMessageRequest(params=params):
with responder:
# Check if this is a task-augmented request
if params.task is not None:
response = await self._task_handlers.augmented_sampling(ctx, params, params.task)
else:
response = await self._sampling_callback(ctx, params)
client_response = ClientResponse.validate_python(response)
await responder.respond(client_response)
case types.ElicitRequest(params=params):
with responder:
# Check if this is a task-augmented request
if params.task is not None:
response = await self._task_handlers.augmented_elicitation(ctx, params, params.task)
else:
response = await self._elicitation_callback(ctx, params)
client_response = ClientResponse.validate_python(response)
await responder.respond(client_response)
case types.ListRootsRequest():
with responder:
response = await self._list_roots_callback(ctx)
client_response = ClientResponse.validate_python(response)
await responder.respond(client_response)
case types.PingRequest(): # pragma: no cover
with responder:
return await responder.respond(types.ClientResult(root=types.EmptyResult()))
case _: # pragma: no cover
pass # Task requests handled above by _task_handlers
return None
async def _handle_incoming(
self,
req: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
) -> None:
"""Handle incoming messages by forwarding to the message handler."""
await self._message_handler(req)
async def _received_notification(self, notification: types.ServerNotification) -> None:
"""Handle notifications from the server."""
# Process specific notification types
match notification.root:
case types.LoggingMessageNotification(params=params):
await self._logging_callback(params)
case types.ElicitCompleteNotification(params=params):
# Handle elicitation completion notification
# Clients MAY use this to retry requests or update UI
# The notification contains the elicitationId of the completed elicitation
pass
case _:
pass
@@ -0,0 +1,447 @@
"""
SessionGroup concurrently manages multiple MCP session connections.
Tools, resources, and prompts are aggregated across servers. Servers may
be connected to or disconnected from at any point after initialization.
This abstractions can handle naming collisions using a custom user-provided
hook.
"""
import contextlib
import logging
from collections.abc import Callable
from dataclasses import dataclass
from datetime import timedelta
from types import TracebackType
from typing import Any, TypeAlias, overload
import anyio
import httpx
from pydantic import BaseModel
from typing_extensions import Self, deprecated
import mcp
from mcp import types
from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
from mcp.client.sse import sse_client
from mcp.client.stdio import StdioServerParameters
from mcp.client.streamable_http import streamable_http_client
from mcp.shared._httpx_utils import create_mcp_http_client
from mcp.shared.exceptions import McpError
from mcp.shared.session import ProgressFnT
class SseServerParameters(BaseModel):
"""Parameters for intializing a sse_client."""
# The endpoint URL.
url: str
# Optional headers to include in requests.
headers: dict[str, Any] | None = None
# HTTP timeout for regular operations.
timeout: float = 5
# Timeout for SSE read operations.
sse_read_timeout: float = 60 * 5
class StreamableHttpParameters(BaseModel):
"""Parameters for intializing a streamable_http_client."""
# The endpoint URL.
url: str
# Optional headers to include in requests.
headers: dict[str, Any] | None = None
# HTTP timeout for regular operations.
timeout: timedelta = timedelta(seconds=30)
# Timeout for SSE read operations.
sse_read_timeout: timedelta = timedelta(seconds=60 * 5)
# Close the client session when the transport closes.
terminate_on_close: bool = True
ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters
# Use dataclass instead of pydantic BaseModel
# because pydantic BaseModel cannot handle Protocol fields.
@dataclass
class ClientSessionParameters:
"""Parameters for establishing a client session to an MCP server."""
read_timeout_seconds: timedelta | None = None
sampling_callback: SamplingFnT | None = None
elicitation_callback: ElicitationFnT | None = None
list_roots_callback: ListRootsFnT | None = None
logging_callback: LoggingFnT | None = None
message_handler: MessageHandlerFnT | None = None
client_info: types.Implementation | None = None
class ClientSessionGroup:
"""Client for managing connections to multiple MCP servers.
This class is responsible for encapsulating management of server connections.
It aggregates tools, resources, and prompts from all connected servers.
For auxiliary handlers, such as resource subscription, this is delegated to
the client and can be accessed via the session.
Example Usage:
name_fn = lambda name, server_info: f"{(server_info.name)}_{name}"
async with ClientSessionGroup(component_name_hook=name_fn) as group:
for server_param in server_params:
await group.connect_to_server(server_param)
...
"""
class _ComponentNames(BaseModel):
"""Used for reverse index to find components."""
prompts: set[str] = set()
resources: set[str] = set()
tools: set[str] = set()
# Standard MCP components.
_prompts: dict[str, types.Prompt]
_resources: dict[str, types.Resource]
_tools: dict[str, types.Tool]
# Client-server connection management.
_sessions: dict[mcp.ClientSession, _ComponentNames]
_tool_to_session: dict[str, mcp.ClientSession]
_exit_stack: contextlib.AsyncExitStack
_session_exit_stacks: dict[mcp.ClientSession, contextlib.AsyncExitStack]
# Optional fn consuming (component_name, serverInfo) for custom names.
# This is provide a means to mitigate naming conflicts across servers.
# Example: (tool_name, serverInfo) => "{result.serverInfo.name}.{tool_name}"
_ComponentNameHook: TypeAlias = Callable[[str, types.Implementation], str]
_component_name_hook: _ComponentNameHook | None
def __init__(
self,
exit_stack: contextlib.AsyncExitStack | None = None,
component_name_hook: _ComponentNameHook | None = None,
) -> None:
"""Initializes the MCP client."""
self._tools = {}
self._resources = {}
self._prompts = {}
self._sessions = {}
self._tool_to_session = {}
if exit_stack is None:
self._exit_stack = contextlib.AsyncExitStack()
self._owns_exit_stack = True
else:
self._exit_stack = exit_stack
self._owns_exit_stack = False
self._session_exit_stacks = {}
self._component_name_hook = component_name_hook
async def __aenter__(self) -> Self: # pragma: no cover
# Enter the exit stack only if we created it ourselves
if self._owns_exit_stack:
await self._exit_stack.__aenter__()
return self
async def __aexit__(
self,
_exc_type: type[BaseException] | None,
_exc_val: BaseException | None,
_exc_tb: TracebackType | None,
) -> bool | None: # pragma: no cover
"""Closes session exit stacks and main exit stack upon completion."""
# Only close the main exit stack if we created it
if self._owns_exit_stack:
await self._exit_stack.aclose()
# Concurrently close session stacks.
async with anyio.create_task_group() as tg:
for exit_stack in self._session_exit_stacks.values():
tg.start_soon(exit_stack.aclose)
@property
def sessions(self) -> list[mcp.ClientSession]:
"""Returns the list of sessions being managed."""
return list(self._sessions.keys()) # pragma: no cover
@property
def prompts(self) -> dict[str, types.Prompt]:
"""Returns the prompts as a dictionary of names to prompts."""
return self._prompts
@property
def resources(self) -> dict[str, types.Resource]:
"""Returns the resources as a dictionary of names to resources."""
return self._resources
@property
def tools(self) -> dict[str, types.Tool]:
"""Returns the tools as a dictionary of names to tools."""
return self._tools
@overload
async def call_tool(
self,
name: str,
arguments: dict[str, Any],
read_timeout_seconds: timedelta | None = None,
progress_callback: ProgressFnT | None = None,
*,
meta: dict[str, Any] | None = None,
) -> types.CallToolResult: ...
@overload
@deprecated("The 'args' parameter is deprecated. Use 'arguments' instead.")
async def call_tool(
self,
name: str,
*,
args: dict[str, Any],
read_timeout_seconds: timedelta | None = None,
progress_callback: ProgressFnT | None = None,
meta: dict[str, Any] | None = None,
) -> types.CallToolResult: ...
async def call_tool(
self,
name: str,
arguments: dict[str, Any] | None = None,
read_timeout_seconds: timedelta | None = None,
progress_callback: ProgressFnT | None = None,
*,
meta: dict[str, Any] | None = None,
args: dict[str, Any] | None = None,
) -> types.CallToolResult:
"""Executes a tool given its name and arguments."""
session = self._tool_to_session[name]
session_tool_name = self.tools[name].name
return await session.call_tool(
session_tool_name,
arguments if args is None else args,
read_timeout_seconds=read_timeout_seconds,
progress_callback=progress_callback,
meta=meta,
)
async def disconnect_from_server(self, session: mcp.ClientSession) -> None:
"""Disconnects from a single MCP server."""
session_known_for_components = session in self._sessions
session_known_for_stack = session in self._session_exit_stacks
if not session_known_for_components and not session_known_for_stack:
raise McpError(
types.ErrorData(
code=types.INVALID_PARAMS,
message="Provided session is not managed or already disconnected.",
)
)
if session_known_for_components: # pragma: no cover
component_names = self._sessions.pop(session) # Pop from _sessions tracking
# Remove prompts associated with the session.
for name in component_names.prompts:
if name in self._prompts:
del self._prompts[name]
# Remove resources associated with the session.
for name in component_names.resources:
if name in self._resources:
del self._resources[name]
# Remove tools associated with the session.
for name in component_names.tools:
if name in self._tools:
del self._tools[name]
if name in self._tool_to_session:
del self._tool_to_session[name]
# Clean up the session's resources via its dedicated exit stack
if session_known_for_stack:
session_stack_to_close = self._session_exit_stacks.pop(session) # pragma: no cover
await session_stack_to_close.aclose() # pragma: no cover
async def connect_with_session(
self, server_info: types.Implementation, session: mcp.ClientSession
) -> mcp.ClientSession:
"""Connects to a single MCP server."""
await self._aggregate_components(server_info, session)
return session
async def connect_to_server(
self,
server_params: ServerParameters,
session_params: ClientSessionParameters | None = None,
) -> mcp.ClientSession:
"""Connects to a single MCP server."""
server_info, session = await self._establish_session(server_params, session_params or ClientSessionParameters())
return await self.connect_with_session(server_info, session)
async def _establish_session(
self,
server_params: ServerParameters,
session_params: ClientSessionParameters,
) -> tuple[types.Implementation, mcp.ClientSession]:
"""Establish a client session to an MCP server."""
session_stack = contextlib.AsyncExitStack()
try:
# Create read and write streams that facilitate io with the server.
if isinstance(server_params, StdioServerParameters):
client = mcp.stdio_client(server_params)
read, write = await session_stack.enter_async_context(client)
elif isinstance(server_params, SseServerParameters):
client = sse_client(
url=server_params.url,
headers=server_params.headers,
timeout=server_params.timeout,
sse_read_timeout=server_params.sse_read_timeout,
)
read, write = await session_stack.enter_async_context(client)
else:
httpx_client = create_mcp_http_client(
headers=server_params.headers,
timeout=httpx.Timeout(
server_params.timeout.total_seconds(),
read=server_params.sse_read_timeout.total_seconds(),
),
)
await session_stack.enter_async_context(httpx_client)
client = streamable_http_client(
url=server_params.url,
http_client=httpx_client,
terminate_on_close=server_params.terminate_on_close,
)
read, write, _ = await session_stack.enter_async_context(client)
session = await session_stack.enter_async_context(
mcp.ClientSession(
read,
write,
read_timeout_seconds=session_params.read_timeout_seconds,
sampling_callback=session_params.sampling_callback,
elicitation_callback=session_params.elicitation_callback,
list_roots_callback=session_params.list_roots_callback,
logging_callback=session_params.logging_callback,
message_handler=session_params.message_handler,
client_info=session_params.client_info,
)
)
result = await session.initialize()
# Session successfully initialized.
# Store its stack and register the stack with the main group stack.
self._session_exit_stacks[session] = session_stack
# session_stack itself becomes a resource managed by the
# main _exit_stack.
await self._exit_stack.enter_async_context(session_stack)
return result.serverInfo, session
except Exception: # pragma: no cover
# If anything during this setup fails, ensure the session-specific
# stack is closed.
await session_stack.aclose()
raise
async def _aggregate_components(self, server_info: types.Implementation, session: mcp.ClientSession) -> None:
"""Aggregates prompts, resources, and tools from a given session."""
# Create a reverse index so we can find all prompts, resources, and
# tools belonging to this session. Used for removing components from
# the session group via self.disconnect_from_server.
component_names = self._ComponentNames()
# Temporary components dicts. We do not want to modify the aggregate
# lists in case of an intermediate failure.
prompts_temp: dict[str, types.Prompt] = {}
resources_temp: dict[str, types.Resource] = {}
tools_temp: dict[str, types.Tool] = {}
tool_to_session_temp: dict[str, mcp.ClientSession] = {}
# Query the server for its prompts and aggregate to list.
try:
prompts = (await session.list_prompts()).prompts
for prompt in prompts:
name = self._component_name(prompt.name, server_info)
prompts_temp[name] = prompt
component_names.prompts.add(name)
except McpError as err: # pragma: no cover
logging.warning(f"Could not fetch prompts: {err}")
# Query the server for its resources and aggregate to list.
try:
resources = (await session.list_resources()).resources
for resource in resources:
name = self._component_name(resource.name, server_info)
resources_temp[name] = resource
component_names.resources.add(name)
except McpError as err: # pragma: no cover
logging.warning(f"Could not fetch resources: {err}")
# Query the server for its tools and aggregate to list.
try:
tools = (await session.list_tools()).tools
for tool in tools:
name = self._component_name(tool.name, server_info)
tools_temp[name] = tool
tool_to_session_temp[name] = session
component_names.tools.add(name)
except McpError as err: # pragma: no cover
logging.warning(f"Could not fetch tools: {err}")
# Clean up exit stack for session if we couldn't retrieve anything
# from the server.
if not any((prompts_temp, resources_temp, tools_temp)):
del self._session_exit_stacks[session] # pragma: no cover
# Check for duplicates.
matching_prompts = prompts_temp.keys() & self._prompts.keys()
if matching_prompts:
raise McpError( # pragma: no cover
types.ErrorData(
code=types.INVALID_PARAMS,
message=f"{matching_prompts} already exist in group prompts.",
)
)
matching_resources = resources_temp.keys() & self._resources.keys()
if matching_resources:
raise McpError( # pragma: no cover
types.ErrorData(
code=types.INVALID_PARAMS,
message=f"{matching_resources} already exist in group resources.",
)
)
matching_tools = tools_temp.keys() & self._tools.keys()
if matching_tools:
raise McpError(
types.ErrorData(
code=types.INVALID_PARAMS,
message=f"{matching_tools} already exist in group tools.",
)
)
# Aggregate components.
self._sessions[session] = component_names
self._prompts.update(prompts_temp)
self._resources.update(resources_temp)
self._tools.update(tools_temp)
self._tool_to_session.update(tool_to_session_temp)
def _component_name(self, name: str, server_info: types.Implementation) -> str:
if self._component_name_hook:
return self._component_name_hook(name, server_info)
return name
@@ -0,0 +1,164 @@
import logging
from collections.abc import Callable
from contextlib import asynccontextmanager
from typing import Any
from urllib.parse import parse_qs, urljoin, urlparse
import anyio
import httpx
from anyio.abc import TaskStatus
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx_sse import aconnect_sse
from httpx_sse._exceptions import SSEError
import mcp.types as types
from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
from mcp.shared.message import SessionMessage
logger = logging.getLogger(__name__)
def remove_request_params(url: str) -> str:
return urljoin(url, urlparse(url).path)
def _extract_session_id_from_endpoint(endpoint_url: str) -> str | None:
query_params = parse_qs(urlparse(endpoint_url).query)
return query_params.get("sessionId", [None])[0] or query_params.get("session_id", [None])[0]
@asynccontextmanager
async def sse_client(
url: str,
headers: dict[str, Any] | None = None,
timeout: float = 5,
sse_read_timeout: float = 60 * 5,
httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
auth: httpx.Auth | None = None,
on_session_created: Callable[[str], None] | None = None,
):
"""
Client transport for SSE.
`sse_read_timeout` determines how long (in seconds) the client will wait for a new
event before disconnecting. All other HTTP operations are controlled by `timeout`.
Args:
url: The SSE endpoint URL.
headers: Optional headers to include in requests.
timeout: HTTP timeout for regular operations.
sse_read_timeout: Timeout for SSE read operations.
auth: Optional HTTPX authentication handler.
on_session_created: Optional callback invoked with the session ID when received.
"""
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
async with anyio.create_task_group() as tg:
try:
logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}")
async with httpx_client_factory(
headers=headers, auth=auth, timeout=httpx.Timeout(timeout, read=sse_read_timeout)
) as client:
async with aconnect_sse(
client,
"GET",
url,
) as event_source:
event_source.response.raise_for_status()
logger.debug("SSE connection established")
async def sse_reader(
task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED,
):
try:
async for sse in event_source.aiter_sse(): # pragma: no branch
logger.debug(f"Received SSE event: {sse.event}")
match sse.event:
case "endpoint":
endpoint_url = urljoin(url, sse.data)
logger.debug(f"Received endpoint URL: {endpoint_url}")
url_parsed = urlparse(url)
endpoint_parsed = urlparse(endpoint_url)
if ( # pragma: no cover
url_parsed.netloc != endpoint_parsed.netloc
or url_parsed.scheme != endpoint_parsed.scheme
):
error_msg = ( # pragma: no cover
f"Endpoint origin does not match connection origin: {endpoint_url}"
)
logger.error(error_msg) # pragma: no cover
raise ValueError(error_msg) # pragma: no cover
if on_session_created:
session_id = _extract_session_id_from_endpoint(endpoint_url)
if session_id:
on_session_created(session_id)
task_status.started(endpoint_url)
case "message":
# Skip empty data (keep-alive pings)
if not sse.data:
continue
try:
message = types.JSONRPCMessage.model_validate_json( # noqa: E501
sse.data
)
logger.debug(f"Received server message: {message}")
except Exception as exc: # pragma: no cover
logger.exception("Error parsing server message") # pragma: no cover
await read_stream_writer.send(exc) # pragma: no cover
continue # pragma: no cover
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
case _: # pragma: no cover
logger.warning(f"Unknown SSE event: {sse.event}") # pragma: no cover
except SSEError as sse_exc: # pragma: no cover
logger.exception("Encountered SSE exception") # pragma: no cover
raise sse_exc # pragma: no cover
except Exception as exc: # pragma: no cover
logger.exception("Error in sse_reader") # pragma: no cover
await read_stream_writer.send(exc) # pragma: no cover
finally:
await read_stream_writer.aclose()
async def post_writer(endpoint_url: str):
try:
async with write_stream_reader:
async for session_message in write_stream_reader:
logger.debug(f"Sending client message: {session_message}")
response = await client.post(
endpoint_url,
json=session_message.message.model_dump(
by_alias=True,
mode="json",
exclude_none=True,
),
)
response.raise_for_status()
logger.debug(f"Client message sent successfully: {response.status_code}")
except Exception: # pragma: no cover
logger.exception("Error in post_writer") # pragma: no cover
finally:
await write_stream.aclose()
endpoint_url = await tg.start(sse_reader)
logger.debug(f"Starting post writer with endpoint URL: {endpoint_url}")
tg.start_soon(post_writer, endpoint_url)
try:
yield read_stream, write_stream
finally:
tg.cancel_scope.cancel()
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
@@ -0,0 +1,278 @@
import logging
import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Literal, TextIO
import anyio
import anyio.lowlevel
from anyio.abc import Process
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from anyio.streams.text import TextReceiveStream
from pydantic import BaseModel, Field
import mcp.types as types
from mcp.os.posix.utilities import terminate_posix_process_tree
from mcp.os.win32.utilities import (
FallbackProcess,
create_windows_process,
get_windows_executable_command,
terminate_windows_process_tree,
)
from mcp.shared.message import SessionMessage
logger = logging.getLogger(__name__)
# Environment variables to inherit by default
DEFAULT_INHERITED_ENV_VARS = (
[
"APPDATA",
"HOMEDRIVE",
"HOMEPATH",
"LOCALAPPDATA",
"PATH",
"PATHEXT",
"PROCESSOR_ARCHITECTURE",
"SYSTEMDRIVE",
"SYSTEMROOT",
"TEMP",
"USERNAME",
"USERPROFILE",
]
if sys.platform == "win32"
else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"]
)
# Timeout for process termination before falling back to force kill
PROCESS_TERMINATION_TIMEOUT = 2.0
def get_default_environment() -> dict[str, str]:
"""
Returns a default environment object including only environment variables deemed
safe to inherit.
"""
env: dict[str, str] = {}
for key in DEFAULT_INHERITED_ENV_VARS:
value = os.environ.get(key)
if value is None:
continue # pragma: no cover
if value.startswith("()"): # pragma: no cover
# Skip functions, which are a security risk
continue # pragma: no cover
env[key] = value
return env
class StdioServerParameters(BaseModel):
command: str
"""The executable to run to start the server."""
args: list[str] = Field(default_factory=list)
"""Command line arguments to pass to the executable."""
env: dict[str, str] | None = None
"""
The environment to use when spawning the process.
If not specified, the result of get_default_environment() will be used.
"""
cwd: str | Path | None = None
"""The working directory to use when spawning the process."""
encoding: str = "utf-8"
"""
The text encoding used when sending/receiving messages to the server
defaults to utf-8
"""
encoding_error_handler: Literal["strict", "ignore", "replace"] = "strict"
"""
The text encoding error handler.
See https://docs.python.org/3/library/codecs.html#codec-base-classes for
explanations of possible values
"""
@asynccontextmanager
async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stderr):
"""
Client transport for stdio: this will connect to a server by spawning a
process and communicating with it over stdin/stdout.
"""
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
try:
command = _get_executable_command(server.command)
# Open process with stderr piped for capture
process = await _create_platform_compatible_process(
command=command,
args=server.args,
env=({**get_default_environment(), **server.env} if server.env is not None else get_default_environment()),
errlog=errlog,
cwd=server.cwd,
)
except OSError:
# Clean up streams if process creation fails
await read_stream.aclose()
await write_stream.aclose()
await read_stream_writer.aclose()
await write_stream_reader.aclose()
raise
async def stdout_reader():
assert process.stdout, "Opened process is missing stdout"
try:
async with read_stream_writer:
buffer = ""
async for chunk in TextReceiveStream(
process.stdout,
encoding=server.encoding,
errors=server.encoding_error_handler,
):
lines = (buffer + chunk).split("\n")
buffer = lines.pop()
for line in lines:
try:
message = types.JSONRPCMessage.model_validate_json(line)
except Exception as exc: # pragma: no cover
logger.exception("Failed to parse JSONRPC message from server")
await read_stream_writer.send(exc)
continue
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
except anyio.ClosedResourceError: # pragma: no cover
await anyio.lowlevel.checkpoint()
async def stdin_writer():
assert process.stdin, "Opened process is missing stdin"
try:
async with write_stream_reader:
async for session_message in write_stream_reader:
json = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
await process.stdin.send(
(json + "\n").encode(
encoding=server.encoding,
errors=server.encoding_error_handler,
)
)
except anyio.ClosedResourceError: # pragma: no cover
await anyio.lowlevel.checkpoint()
async with (
anyio.create_task_group() as tg,
process,
):
tg.start_soon(stdout_reader)
tg.start_soon(stdin_writer)
try:
yield read_stream, write_stream
finally:
# MCP spec: stdio shutdown sequence
# 1. Close input stream to server
# 2. Wait for server to exit, or send SIGTERM if it doesn't exit in time
# 3. Send SIGKILL if still not exited
if process.stdin: # pragma: no branch
try:
await process.stdin.aclose()
except Exception: # pragma: no cover
# stdin might already be closed, which is fine
pass
try:
# Give the process time to exit gracefully after stdin closes
with anyio.fail_after(PROCESS_TERMINATION_TIMEOUT):
await process.wait()
except TimeoutError:
# Process didn't exit from stdin closure, use platform-specific termination
# which handles SIGTERM -> SIGKILL escalation
await _terminate_process_tree(process)
except ProcessLookupError: # pragma: no cover
# Process already exited, which is fine
pass
await read_stream.aclose()
await write_stream.aclose()
await read_stream_writer.aclose()
await write_stream_reader.aclose()
def _get_executable_command(command: str) -> str:
"""
Get the correct executable command normalized for the current platform.
Args:
command: Base command (e.g., 'uvx', 'npx')
Returns:
str: Platform-appropriate command
"""
if sys.platform == "win32": # pragma: no cover
return get_windows_executable_command(command)
else:
return command # pragma: no cover
async def _create_platform_compatible_process(
command: str,
args: list[str],
env: dict[str, str] | None = None,
errlog: TextIO = sys.stderr,
cwd: Path | str | None = None,
):
"""
Creates a subprocess in a platform-compatible way.
Unix: Creates process in a new session/process group for killpg support
Windows: Creates process in a Job Object for reliable child termination
"""
if sys.platform == "win32": # pragma: no cover
process = await create_windows_process(command, args, env, errlog, cwd)
else:
process = await anyio.open_process(
[command, *args],
env=env,
stderr=errlog,
cwd=cwd,
start_new_session=True,
) # pragma: no cover
return process
async def _terminate_process_tree(process: Process | FallbackProcess, timeout_seconds: float = 2.0) -> None:
"""
Terminate a process and all its children using platform-specific methods.
Unix: Uses os.killpg() for atomic process group termination
Windows: Uses Job Objects via pywin32 for reliable child process cleanup
Args:
process: The process to terminate
timeout_seconds: Timeout in seconds before force killing (default: 2.0)
"""
if sys.platform == "win32": # pragma: no cover
await terminate_windows_process_tree(process, timeout_seconds)
else: # pragma: no cover
# FallbackProcess should only be used for Windows compatibility
assert isinstance(process, Process)
await terminate_posix_process_tree(process, timeout_seconds)
@@ -0,0 +1,722 @@
"""
StreamableHTTP Client Transport Module
This module implements the StreamableHTTP transport for MCP clients,
providing support for HTTP POST requests with optional SSE streaming responses
and session management.
"""
import contextlib
import logging
from collections.abc import AsyncGenerator, Awaitable, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, overload
from warnings import warn
import anyio
import httpx
from anyio.abc import TaskGroup
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
from typing_extensions import deprecated
from mcp.shared._httpx_utils import (
McpHttpClientFactory,
create_mcp_http_client,
)
from mcp.shared.message import ClientMessageMetadata, SessionMessage
from mcp.types import (
ErrorData,
InitializeResult,
JSONRPCError,
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
JSONRPCResponse,
RequestId,
)
logger = logging.getLogger(__name__)
SessionMessageOrError = SessionMessage | Exception
StreamWriter = MemoryObjectSendStream[SessionMessageOrError]
StreamReader = MemoryObjectReceiveStream[SessionMessage]
GetSessionIdCallback = Callable[[], str | None]
MCP_SESSION_ID = "mcp-session-id"
MCP_PROTOCOL_VERSION = "mcp-protocol-version"
LAST_EVENT_ID = "last-event-id"
# Reconnection defaults
DEFAULT_RECONNECTION_DELAY_MS = 1000 # 1 second fallback when server doesn't provide retry
MAX_RECONNECTION_ATTEMPTS = 2 # Max retry attempts before giving up
CONTENT_TYPE = "content-type"
ACCEPT = "accept"
JSON = "application/json"
SSE = "text/event-stream"
# Sentinel value for detecting unset optional parameters
_UNSET = object()
class StreamableHTTPError(Exception):
"""Base exception for StreamableHTTP transport errors."""
class ResumptionError(StreamableHTTPError):
"""Raised when resumption request is invalid."""
@dataclass
class RequestContext:
"""Context for a request operation."""
client: httpx.AsyncClient
session_id: str | None
session_message: SessionMessage
metadata: ClientMessageMetadata | None
read_stream_writer: StreamWriter
headers: dict[str, str] | None = None # Deprecated - no longer used
sse_read_timeout: float | None = None # Deprecated - no longer used
class StreamableHTTPTransport:
"""StreamableHTTP client transport implementation."""
@overload
def __init__(self, url: str) -> None: ...
@overload
@deprecated(
"Parameters headers, timeout, sse_read_timeout, and auth are deprecated. "
"Configure these on the httpx.AsyncClient instead."
)
def __init__(
self,
url: str,
headers: dict[str, str] | None = None,
timeout: float | timedelta = 30,
sse_read_timeout: float | timedelta = 60 * 5,
auth: httpx.Auth | None = None,
) -> None: ...
def __init__(
self,
url: str,
headers: Any = _UNSET,
timeout: Any = _UNSET,
sse_read_timeout: Any = _UNSET,
auth: Any = _UNSET,
) -> None:
"""Initialize the StreamableHTTP transport.
Args:
url: The endpoint URL.
headers: Optional headers to include in requests.
timeout: HTTP timeout for regular operations.
sse_read_timeout: Timeout for SSE read operations.
auth: Optional HTTPX authentication handler.
"""
# Check for deprecated parameters and issue runtime warning
deprecated_params: list[str] = []
if headers is not _UNSET:
deprecated_params.append("headers")
if timeout is not _UNSET:
deprecated_params.append("timeout")
if sse_read_timeout is not _UNSET:
deprecated_params.append("sse_read_timeout")
if auth is not _UNSET:
deprecated_params.append("auth")
if deprecated_params:
warn(
f"Parameters {', '.join(deprecated_params)} are deprecated and will be ignored. "
"Configure these on the httpx.AsyncClient instead.",
DeprecationWarning,
stacklevel=2,
)
self.url = url
self.session_id = None
self.protocol_version = None
def _prepare_headers(self) -> dict[str, str]:
"""Build MCP-specific request headers.
These headers will be merged with the httpx.AsyncClient's default headers,
with these MCP-specific headers taking precedence.
"""
headers: dict[str, str] = {}
# Add MCP protocol headers
headers[ACCEPT] = f"{JSON}, {SSE}"
headers[CONTENT_TYPE] = JSON
# Add session headers if available
if self.session_id:
headers[MCP_SESSION_ID] = self.session_id
if self.protocol_version:
headers[MCP_PROTOCOL_VERSION] = self.protocol_version
return headers
def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
"""Check if the message is an initialization request."""
return isinstance(message.root, JSONRPCRequest) and message.root.method == "initialize"
def _is_initialized_notification(self, message: JSONRPCMessage) -> bool:
"""Check if the message is an initialized notification."""
return isinstance(message.root, JSONRPCNotification) and message.root.method == "notifications/initialized"
def _maybe_extract_session_id_from_response(
self,
response: httpx.Response,
) -> None:
"""Extract and store session ID from response headers."""
new_session_id = response.headers.get(MCP_SESSION_ID)
if new_session_id:
self.session_id = new_session_id
logger.info(f"Received session ID: {self.session_id}")
def _maybe_extract_protocol_version_from_message(
self,
message: JSONRPCMessage,
) -> None:
"""Extract protocol version from initialization response message."""
if isinstance(message.root, JSONRPCResponse) and message.root.result: # pragma: no branch
try:
# Parse the result as InitializeResult for type safety
init_result = InitializeResult.model_validate(message.root.result)
self.protocol_version = str(init_result.protocolVersion)
logger.info(f"Negotiated protocol version: {self.protocol_version}")
except Exception as exc: # pragma: no cover
logger.warning(
f"Failed to parse initialization response as InitializeResult: {exc}"
) # pragma: no cover
logger.warning(f"Raw result: {message.root.result}")
async def _handle_sse_event(
self,
sse: ServerSentEvent,
read_stream_writer: StreamWriter,
original_request_id: RequestId | None = None,
resumption_callback: Callable[[str], Awaitable[None]] | None = None,
is_initialization: bool = False,
) -> bool:
"""Handle an SSE event, returning True if the response is complete."""
if sse.event == "message":
# Handle priming events (empty data with ID) for resumability
if not sse.data:
# Call resumption callback for priming events that have an ID
if sse.id and resumption_callback:
await resumption_callback(sse.id)
return False
try:
message = JSONRPCMessage.model_validate_json(sse.data)
logger.debug(f"SSE message: {message}")
# Extract protocol version from initialization response
if is_initialization:
self._maybe_extract_protocol_version_from_message(message)
# If this is a response and we have original_request_id, replace it
if original_request_id is not None and isinstance(message.root, JSONRPCResponse | JSONRPCError):
message.root.id = original_request_id
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
# Call resumption token callback if we have an ID
if sse.id and resumption_callback:
await resumption_callback(sse.id)
# If this is a response or error return True indicating completion
# Otherwise, return False to continue listening
return isinstance(message.root, JSONRPCResponse | JSONRPCError)
except Exception as exc: # pragma: no cover
logger.exception("Error parsing SSE message")
await read_stream_writer.send(exc)
return False
else: # pragma: no cover
logger.warning(f"Unknown SSE event: {sse.event}")
return False
async def handle_get_stream(
self,
client: httpx.AsyncClient,
read_stream_writer: StreamWriter,
) -> None:
"""Handle GET stream for server-initiated messages with auto-reconnect."""
last_event_id: str | None = None
retry_interval_ms: int | None = None
attempt: int = 0
while attempt < MAX_RECONNECTION_ATTEMPTS: # pragma: no branch
try:
if not self.session_id:
return
headers = self._prepare_headers()
if last_event_id:
headers[LAST_EVENT_ID] = last_event_id # pragma: no cover
async with aconnect_sse(
client,
"GET",
self.url,
headers=headers,
) as event_source:
event_source.response.raise_for_status()
logger.debug("GET SSE connection established")
async for sse in event_source.aiter_sse():
# Track last event ID for reconnection
if sse.id:
last_event_id = sse.id # pragma: no cover
# Track retry interval from server
if sse.retry is not None:
retry_interval_ms = sse.retry # pragma: no cover
await self._handle_sse_event(sse, read_stream_writer)
# Stream ended normally (server closed) - reset attempt counter
attempt = 0
except Exception as exc: # pragma: no cover
logger.debug(f"GET stream error: {exc}")
attempt += 1
if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover
logger.debug(f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded")
return
# Wait before reconnecting
delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS
logger.info(f"GET stream disconnected, reconnecting in {delay_ms}ms...")
await anyio.sleep(delay_ms / 1000.0)
async def _handle_resumption_request(self, ctx: RequestContext) -> None:
"""Handle a resumption request using GET with SSE."""
headers = self._prepare_headers()
if ctx.metadata and ctx.metadata.resumption_token:
headers[LAST_EVENT_ID] = ctx.metadata.resumption_token
else:
raise ResumptionError("Resumption request requires a resumption token") # pragma: no cover
# Extract original request ID to map responses
original_request_id = None
if isinstance(ctx.session_message.message.root, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.root.id
async with aconnect_sse(
ctx.client,
"GET",
self.url,
headers=headers,
) as event_source:
event_source.response.raise_for_status()
logger.debug("Resumption GET SSE connection established")
async for sse in event_source.aiter_sse(): # pragma: no branch
is_complete = await self._handle_sse_event(
sse,
ctx.read_stream_writer,
original_request_id,
ctx.metadata.on_resumption_token_update if ctx.metadata else None,
)
if is_complete:
await event_source.response.aclose()
break
async def _handle_post_request(self, ctx: RequestContext) -> None:
"""Handle a POST request with response processing."""
headers = self._prepare_headers()
message = ctx.session_message.message
is_initialization = self._is_initialization_request(message)
async with ctx.client.stream(
"POST",
self.url,
json=message.model_dump(by_alias=True, mode="json", exclude_none=True),
headers=headers,
) as response:
if response.status_code == 202:
logger.debug("Received 202 Accepted")
return
if response.status_code == 404: # pragma: no branch
if isinstance(message.root, JSONRPCRequest):
await self._send_session_terminated_error( # pragma: no cover
ctx.read_stream_writer, # pragma: no cover
message.root.id, # pragma: no cover
) # pragma: no cover
return # pragma: no cover
response.raise_for_status()
if is_initialization:
self._maybe_extract_session_id_from_response(response)
# Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications:
# The server MUST NOT send a response to notifications.
if isinstance(message.root, JSONRPCRequest):
content_type = response.headers.get(CONTENT_TYPE, "").lower()
if content_type.startswith(JSON):
await self._handle_json_response(response, ctx.read_stream_writer, is_initialization)
elif content_type.startswith(SSE):
await self._handle_sse_response(response, ctx, is_initialization)
else:
await self._handle_unexpected_content_type( # pragma: no cover
content_type, # pragma: no cover
ctx.read_stream_writer, # pragma: no cover
) # pragma: no cover
async def _handle_json_response(
self,
response: httpx.Response,
read_stream_writer: StreamWriter,
is_initialization: bool = False,
) -> None:
"""Handle JSON response from the server."""
try:
content = await response.aread()
message = JSONRPCMessage.model_validate_json(content)
# Extract protocol version from initialization response
if is_initialization:
self._maybe_extract_protocol_version_from_message(message)
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
except Exception as exc: # pragma: no cover
logger.exception("Error parsing JSON response")
await read_stream_writer.send(exc)
async def _handle_sse_response(
self,
response: httpx.Response,
ctx: RequestContext,
is_initialization: bool = False,
) -> None:
"""Handle SSE response from the server."""
last_event_id: str | None = None
retry_interval_ms: int | None = None
try:
event_source = EventSource(response)
async for sse in event_source.aiter_sse(): # pragma: no branch
# Track last event ID for potential reconnection
if sse.id:
last_event_id = sse.id
# Track retry interval from server
if sse.retry is not None:
retry_interval_ms = sse.retry
is_complete = await self._handle_sse_event(
sse,
ctx.read_stream_writer,
resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None),
is_initialization=is_initialization,
)
# If the SSE event indicates completion, like returning respose/error
# break the loop
if is_complete:
await response.aclose()
return # Normal completion, no reconnect needed
except Exception as e: # pragma: no cover
logger.debug(f"SSE stream ended: {e}")
# Stream ended without response - reconnect if we received an event with ID
if last_event_id is not None: # pragma: no branch
logger.info("SSE stream disconnected, reconnecting...")
await self._handle_reconnection(ctx, last_event_id, retry_interval_ms)
async def _handle_reconnection(
self,
ctx: RequestContext,
last_event_id: str,
retry_interval_ms: int | None = None,
attempt: int = 0,
) -> None:
"""Reconnect with Last-Event-ID to resume stream after server disconnect."""
# Bail if max retries exceeded
if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover
logger.debug(f"Max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded")
return
# Always wait - use server value or default
delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS
await anyio.sleep(delay_ms / 1000.0)
headers = self._prepare_headers()
headers[LAST_EVENT_ID] = last_event_id
# Extract original request ID to map responses
original_request_id = None
if isinstance(ctx.session_message.message.root, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.root.id
try:
async with aconnect_sse(
ctx.client,
"GET",
self.url,
headers=headers,
) as event_source:
event_source.response.raise_for_status()
logger.info("Reconnected to SSE stream")
# Track for potential further reconnection
reconnect_last_event_id: str = last_event_id
reconnect_retry_ms = retry_interval_ms
async for sse in event_source.aiter_sse():
if sse.id: # pragma: no branch
reconnect_last_event_id = sse.id
if sse.retry is not None:
reconnect_retry_ms = sse.retry
is_complete = await self._handle_sse_event(
sse,
ctx.read_stream_writer,
original_request_id,
ctx.metadata.on_resumption_token_update if ctx.metadata else None,
)
if is_complete:
await event_source.response.aclose()
return
# Stream ended again without response - reconnect again (reset attempt counter)
logger.info("SSE stream disconnected, reconnecting...")
await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0)
except Exception as e: # pragma: no cover
logger.debug(f"Reconnection failed: {e}")
# Try to reconnect again if we still have an event ID
await self._handle_reconnection(ctx, last_event_id, retry_interval_ms, attempt + 1)
async def _handle_unexpected_content_type(
self,
content_type: str,
read_stream_writer: StreamWriter,
) -> None: # pragma: no cover
"""Handle unexpected content type in response."""
error_msg = f"Unexpected content type: {content_type}" # pragma: no cover
logger.error(error_msg) # pragma: no cover
await read_stream_writer.send(ValueError(error_msg)) # pragma: no cover
async def _send_session_terminated_error(
self,
read_stream_writer: StreamWriter,
request_id: RequestId,
) -> None:
"""Send a session terminated error response."""
jsonrpc_error = JSONRPCError(
jsonrpc="2.0",
id=request_id,
error=ErrorData(code=32600, message="Session terminated"),
)
session_message = SessionMessage(JSONRPCMessage(jsonrpc_error))
await read_stream_writer.send(session_message)
async def post_writer(
self,
client: httpx.AsyncClient,
write_stream_reader: StreamReader,
read_stream_writer: StreamWriter,
write_stream: MemoryObjectSendStream[SessionMessage],
start_get_stream: Callable[[], None],
tg: TaskGroup,
) -> None:
"""Handle writing requests to the server."""
try:
async with write_stream_reader:
async for session_message in write_stream_reader:
message = session_message.message
metadata = (
session_message.metadata
if isinstance(session_message.metadata, ClientMessageMetadata)
else None
)
# Check if this is a resumption request
is_resumption = bool(metadata and metadata.resumption_token)
logger.debug(f"Sending client message: {message}")
# Handle initialized notification
if self._is_initialized_notification(message):
start_get_stream()
ctx = RequestContext(
client=client,
session_id=self.session_id,
session_message=session_message,
metadata=metadata,
read_stream_writer=read_stream_writer,
)
async def handle_request_async():
if is_resumption:
await self._handle_resumption_request(ctx)
else:
await self._handle_post_request(ctx)
# If this is a request, start a new task to handle it
if isinstance(message.root, JSONRPCRequest):
tg.start_soon(handle_request_async)
else:
await handle_request_async()
except Exception:
logger.exception("Error in post_writer") # pragma: no cover
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
async def terminate_session(self, client: httpx.AsyncClient) -> None: # pragma: no cover
"""Terminate the session by sending a DELETE request."""
if not self.session_id:
return
try:
headers = self._prepare_headers()
response = await client.delete(self.url, headers=headers)
if response.status_code == 405:
logger.debug("Server does not allow session termination")
elif response.status_code not in (200, 204):
logger.warning(f"Session termination failed: {response.status_code}")
except Exception as exc:
logger.warning(f"Session termination failed: {exc}")
def get_session_id(self) -> str | None:
"""Get the current session ID."""
return self.session_id
@asynccontextmanager
async def streamable_http_client(
url: str,
*,
http_client: httpx.AsyncClient | None = None,
terminate_on_close: bool = True,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
GetSessionIdCallback,
],
None,
]:
"""
Client transport for StreamableHTTP.
Args:
url: The MCP server endpoint URL.
http_client: Optional pre-configured httpx.AsyncClient. If None, a default
client with recommended MCP timeouts will be created. To configure headers,
authentication, or other HTTP settings, create an httpx.AsyncClient and pass it here.
terminate_on_close: If True, send a DELETE request to terminate the session
when the context exits.
Yields:
Tuple containing:
- read_stream: Stream for reading messages from the server
- write_stream: Stream for sending messages to the server
- get_session_id_callback: Function to retrieve the current session ID
Example:
See examples/snippets/clients/ for usage patterns.
"""
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
# Determine if we need to create and manage the client
client_provided = http_client is not None
client = http_client
if client is None:
# Create default client with recommended MCP timeouts
client = create_mcp_http_client()
transport = StreamableHTTPTransport(url)
async with anyio.create_task_group() as tg:
try:
logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
async with contextlib.AsyncExitStack() as stack:
# Only manage client lifecycle if we created it
if not client_provided:
await stack.enter_async_context(client)
def start_get_stream() -> None:
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
tg.start_soon(
transport.post_writer,
client,
write_stream_reader,
read_stream_writer,
write_stream,
start_get_stream,
tg,
)
try:
yield (
read_stream,
write_stream,
transport.get_session_id,
)
finally:
if transport.session_id and terminate_on_close:
await transport.terminate_session(client)
tg.cancel_scope.cancel()
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
@asynccontextmanager
@deprecated("Use `streamable_http_client` instead.")
async def streamablehttp_client(
url: str,
headers: dict[str, str] | None = None,
timeout: float | timedelta = 30,
sse_read_timeout: float | timedelta = 60 * 5,
terminate_on_close: bool = True,
httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
auth: httpx.Auth | None = None,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
GetSessionIdCallback,
],
None,
]:
# Convert timeout parameters
timeout_seconds = timeout.total_seconds() if isinstance(timeout, timedelta) else timeout
sse_read_timeout_seconds = (
sse_read_timeout.total_seconds() if isinstance(sse_read_timeout, timedelta) else sse_read_timeout
)
# Create httpx client using the factory with old-style parameters
client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(timeout_seconds, read=sse_read_timeout_seconds),
auth=auth,
)
# Manage client lifecycle since we created it
async with client:
async with streamable_http_client(
url,
http_client=client,
terminate_on_close=terminate_on_close,
) as streams:
yield streams
@@ -0,0 +1,86 @@
import json
import logging
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import ValidationError
from websockets.asyncio.client import connect as ws_connect
from websockets.typing import Subprotocol
import mcp.types as types
from mcp.shared.message import SessionMessage
logger = logging.getLogger(__name__)
@asynccontextmanager
async def websocket_client(
url: str,
) -> AsyncGenerator[
tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]],
None,
]:
"""
WebSocket client transport for MCP, symmetrical to the server version.
Connects to 'url' using the 'mcp' subprotocol, then yields:
(read_stream, write_stream)
- read_stream: As you read from this stream, you'll receive either valid
JSONRPCMessage objects or Exception objects (when validation fails).
- write_stream: Write JSONRPCMessage objects to this stream to send them
over the WebSocket to the server.
"""
# Create two in-memory streams:
# - One for incoming messages (read_stream, written by ws_reader)
# - One for outgoing messages (write_stream, read by ws_writer)
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
# Connect using websockets, requesting the "mcp" subprotocol
async with ws_connect(url, subprotocols=[Subprotocol("mcp")]) as ws:
async def ws_reader():
"""
Reads text messages from the WebSocket, parses them as JSON-RPC messages,
and sends them into read_stream_writer.
"""
async with read_stream_writer:
async for raw_text in ws:
try:
message = types.JSONRPCMessage.model_validate_json(raw_text)
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
except ValidationError as exc: # pragma: no cover
# If JSON parse or model validation fails, send the exception
await read_stream_writer.send(exc)
async def ws_writer():
"""
Reads JSON-RPC messages from write_stream_reader and
sends them to the server.
"""
async with write_stream_reader:
async for session_message in write_stream_reader:
# Convert to a dict, then to JSON
msg_dict = session_message.message.model_dump(by_alias=True, mode="json", exclude_none=True)
await ws.send(json.dumps(msg_dict))
async with anyio.create_task_group() as tg:
# Start reader and writer tasks
tg.start_soon(ws_reader)
tg.start_soon(ws_writer)
# Yield the receive/send streams
yield (read_stream, write_stream)
# Once the caller's 'async with' block exits, we shut down
tg.cancel_scope.cancel()
@@ -0,0 +1 @@
"""Platform-specific utilities for MCP."""
@@ -0,0 +1 @@
"""POSIX-specific utilities for MCP."""
@@ -0,0 +1,60 @@
"""
POSIX-specific functionality for stdio client operations.
"""
import logging
import os
import signal
import anyio
from anyio.abc import Process
logger = logging.getLogger(__name__)
async def terminate_posix_process_tree(process: Process, timeout_seconds: float = 2.0) -> None:
"""
Terminate a process and all its children on POSIX systems.
Uses os.killpg() for atomic process group termination.
Args:
process: The process to terminate
timeout_seconds: Timeout in seconds before force killing (default: 2.0)
"""
pid = getattr(process, "pid", None) or getattr(getattr(process, "popen", None), "pid", None)
if not pid:
# No PID means there's no process to terminate - it either never started,
# already exited, or we have an invalid process object
return
try:
pgid = os.getpgid(pid)
os.killpg(pgid, signal.SIGTERM)
with anyio.move_on_after(timeout_seconds):
while True:
try:
# Check if process group still exists (signal 0 = check only)
os.killpg(pgid, 0)
await anyio.sleep(0.1)
except ProcessLookupError:
return
try:
os.killpg(pgid, signal.SIGKILL)
except ProcessLookupError:
pass
except (ProcessLookupError, PermissionError, OSError) as e:
logger.warning(f"Process group termination failed for PID {pid}: {e}, falling back to simple terminate")
try:
process.terminate()
with anyio.fail_after(timeout_seconds):
await process.wait()
except Exception:
logger.warning(f"Process termination failed for PID {pid}, attempting force kill")
try:
process.kill()
except Exception:
logger.exception(f"Failed to kill process {pid}")
@@ -0,0 +1 @@
"""Windows-specific utilities for MCP."""
@@ -0,0 +1,338 @@
"""
Windows-specific functionality for stdio client operations.
"""
import logging
import shutil
import subprocess
import sys
from pathlib import Path
from typing import BinaryIO, TextIO, cast
import anyio
from anyio import to_thread
from anyio.abc import Process
from anyio.streams.file import FileReadStream, FileWriteStream
from typing_extensions import deprecated
logger = logging.getLogger("client.stdio.win32")
# Windows-specific imports for Job Objects
if sys.platform == "win32":
import pywintypes
import win32api
import win32con
import win32job
else:
# Type stubs for non-Windows platforms
win32api = None
win32con = None
win32job = None
pywintypes = None
JobHandle = int
def get_windows_executable_command(command: str) -> str:
"""
Get the correct executable command normalized for Windows.
On Windows, commands might exist with specific extensions (.exe, .cmd, etc.)
that need to be located for proper execution.
Args:
command: Base command (e.g., 'uvx', 'npx')
Returns:
str: Windows-appropriate command path
"""
try:
# First check if command exists in PATH as-is
if command_path := shutil.which(command):
return command_path
# Check for Windows-specific extensions
for ext in [".cmd", ".bat", ".exe", ".ps1"]:
ext_version = f"{command}{ext}"
if ext_path := shutil.which(ext_version):
return ext_path
# For regular commands or if we couldn't find special versions
return command
except OSError:
# Handle file system errors during path resolution
# (permissions, broken symlinks, etc.)
return command
class FallbackProcess:
"""
A fallback process wrapper for Windows to handle async I/O
when using subprocess.Popen, which provides sync-only FileIO objects.
This wraps stdin and stdout into async-compatible
streams (FileReadStream, FileWriteStream),
so that MCP clients expecting async streams can work properly.
"""
def __init__(self, popen_obj: subprocess.Popen[bytes]):
self.popen: subprocess.Popen[bytes] = popen_obj
self.stdin_raw = popen_obj.stdin # type: ignore[assignment]
self.stdout_raw = popen_obj.stdout # type: ignore[assignment]
self.stderr = popen_obj.stderr # type: ignore[assignment]
self.stdin = FileWriteStream(cast(BinaryIO, self.stdin_raw)) if self.stdin_raw else None
self.stdout = FileReadStream(cast(BinaryIO, self.stdout_raw)) if self.stdout_raw else None
async def __aenter__(self):
"""Support async context manager entry."""
return self
async def __aexit__(
self,
exc_type: BaseException | None,
exc_val: BaseException | None,
exc_tb: object | None,
) -> None:
"""Terminate and wait on process exit inside a thread."""
self.popen.terminate()
await to_thread.run_sync(self.popen.wait)
# Close the file handles to prevent ResourceWarning
if self.stdin:
await self.stdin.aclose()
if self.stdout:
await self.stdout.aclose()
if self.stdin_raw:
self.stdin_raw.close()
if self.stdout_raw:
self.stdout_raw.close()
if self.stderr:
self.stderr.close()
async def wait(self):
"""Async wait for process completion."""
return await to_thread.run_sync(self.popen.wait)
def terminate(self):
"""Terminate the subprocess immediately."""
return self.popen.terminate()
def kill(self) -> None:
"""Kill the subprocess immediately (alias for terminate)."""
self.terminate()
@property
def pid(self) -> int:
"""Return the process ID."""
return self.popen.pid
# ------------------------
# Updated function
# ------------------------
async def create_windows_process(
command: str,
args: list[str],
env: dict[str, str] | None = None,
errlog: TextIO | None = sys.stderr,
cwd: Path | str | None = None,
) -> Process | FallbackProcess:
"""
Creates a subprocess in a Windows-compatible way with Job Object support.
Attempt to use anyio's open_process for async subprocess creation.
In some cases this will throw NotImplementedError on Windows, e.g.
when using the SelectorEventLoop which does not support async subprocesses.
In that case, we fall back to using subprocess.Popen.
The process is automatically added to a Job Object to ensure all child
processes are terminated when the parent is terminated.
Args:
command (str): The executable to run
args (list[str]): List of command line arguments
env (dict[str, str] | None): Environment variables
errlog (TextIO | None): Where to send stderr output (defaults to sys.stderr)
cwd (Path | str | None): Working directory for the subprocess
Returns:
Process | FallbackProcess: Async-compatible subprocess with stdin and stdout streams
"""
job = _create_job_object()
process = None
try:
# First try using anyio with Windows-specific flags to hide console window
process = await anyio.open_process(
[command, *args],
env=env,
# Ensure we don't create console windows for each process
creationflags=subprocess.CREATE_NO_WINDOW # type: ignore
if hasattr(subprocess, "CREATE_NO_WINDOW")
else 0,
stderr=errlog,
cwd=cwd,
)
except NotImplementedError:
# If Windows doesn't support async subprocess creation, use fallback
process = await _create_windows_fallback_process(command, args, env, errlog, cwd)
except Exception:
# Try again without creation flags
process = await anyio.open_process(
[command, *args],
env=env,
stderr=errlog,
cwd=cwd,
)
_maybe_assign_process_to_job(process, job)
return process
async def _create_windows_fallback_process(
command: str,
args: list[str],
env: dict[str, str] | None = None,
errlog: TextIO | None = sys.stderr,
cwd: Path | str | None = None,
) -> FallbackProcess:
"""
Create a subprocess using subprocess.Popen as a fallback when anyio fails.
This function wraps the sync subprocess.Popen in an async-compatible interface.
"""
try:
# Try launching with creationflags to avoid opening a new console window
popen_obj = subprocess.Popen(
[command, *args],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=errlog,
env=env,
cwd=cwd,
bufsize=0, # Unbuffered output
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
)
except Exception:
# If creationflags failed, fallback without them
popen_obj = subprocess.Popen(
[command, *args],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=errlog,
env=env,
cwd=cwd,
bufsize=0,
)
return FallbackProcess(popen_obj)
def _create_job_object() -> int | None:
"""
Create a Windows Job Object configured to terminate all processes when closed.
"""
if sys.platform != "win32" or not win32job:
return None
try:
job = win32job.CreateJobObject(None, "")
extended_info = win32job.QueryInformationJobObject(job, win32job.JobObjectExtendedLimitInformation)
extended_info["BasicLimitInformation"]["LimitFlags"] |= win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
win32job.SetInformationJobObject(job, win32job.JobObjectExtendedLimitInformation, extended_info)
return job
except Exception as e:
logger.warning(f"Failed to create Job Object for process tree management: {e}")
return None
def _maybe_assign_process_to_job(process: Process | FallbackProcess, job: JobHandle | None) -> None:
"""
Try to assign a process to a job object. If assignment fails
for any reason, the job handle is closed.
"""
if not job:
return
if sys.platform != "win32" or not win32api or not win32con or not win32job:
return
try:
process_handle = win32api.OpenProcess(
win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid
)
if not process_handle:
raise Exception("Failed to open process handle")
try:
win32job.AssignProcessToJobObject(job, process_handle)
process._job_object = job
finally:
win32api.CloseHandle(process_handle)
except Exception as e:
logger.warning(f"Failed to assign process {process.pid} to Job Object: {e}")
if win32api:
win32api.CloseHandle(job)
async def terminate_windows_process_tree(process: Process | FallbackProcess, timeout_seconds: float = 2.0) -> None:
"""
Terminate a process and all its children on Windows.
If the process has an associated job object, it will be terminated.
Otherwise, falls back to basic process termination.
Args:
process: The process to terminate
timeout_seconds: Timeout in seconds before force killing (default: 2.0)
"""
if sys.platform != "win32":
return
job = getattr(process, "_job_object", None)
if job and win32job:
try:
win32job.TerminateJobObject(job, 1)
except Exception:
# Job might already be terminated
pass
finally:
if win32api:
try:
win32api.CloseHandle(job)
except Exception:
pass
# Always try to terminate the process itself as well
try:
process.terminate()
except Exception:
pass
@deprecated(
"terminate_windows_process is deprecated and will be removed in a future version. "
"Process termination is now handled internally by the stdio_client context manager."
)
async def terminate_windows_process(process: Process | FallbackProcess):
"""
Terminate a Windows process.
Note: On Windows, terminating a process with process.terminate() doesn't
always guarantee immediate process termination.
So we give it 2s to exit, or we call process.kill()
which sends a SIGKILL equivalent signal.
Args:
process: The process to terminate
"""
try:
process.terminate()
with anyio.fail_after(2.0):
await process.wait()
except TimeoutError:
# Force kill if it doesn't terminate
process.kill()
Whitespace-only changes.
@@ -0,0 +1,5 @@
from .fastmcp import FastMCP
from .lowlevel import NotificationOptions, Server
from .models import InitializationOptions
__all__ = ["Server", "FastMCP", "NotificationOptions", "InitializationOptions"]
@@ -0,0 +1,50 @@
import importlib.metadata
import logging
import sys
import anyio
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.server.stdio import stdio_server
from mcp.types import ServerCapabilities
if not sys.warnoptions:
import warnings
warnings.simplefilter("ignore")
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("server")
async def receive_loop(session: ServerSession):
logger.info("Starting receive loop")
async for message in session.incoming_messages:
if isinstance(message, Exception):
logger.error("Error: %s", message)
continue
logger.info("Received message from client: %s", message)
async def main():
version = importlib.metadata.version("mcp")
async with stdio_server() as (read_stream, write_stream):
async with (
ServerSession(
read_stream,
write_stream,
InitializationOptions(
server_name="mcp",
server_version=version,
capabilities=ServerCapabilities(),
),
) as session,
write_stream,
):
await receive_loop(session)
if __name__ == "__main__":
anyio.run(main, backend="trio")
@@ -0,0 +1,3 @@
"""
MCP OAuth server authorization components.
"""
@@ -0,0 +1,5 @@
from pydantic import ValidationError
def stringify_pydantic_error(validation_error: ValidationError) -> str:
return "\n".join(f"{'.'.join(str(loc) for loc in e['loc'])}: {e['msg']}" for e in validation_error.errors())
@@ -0,0 +1,3 @@
"""
Request handlers for MCP authorization endpoints.
"""
@@ -0,0 +1,224 @@
import logging
from dataclasses import dataclass
from typing import Any, Literal
from pydantic import AnyUrl, BaseModel, Field, RootModel, ValidationError
from starlette.datastructures import FormData, QueryParams
from starlette.requests import Request
from starlette.responses import RedirectResponse, Response
from mcp.server.auth.errors import stringify_pydantic_error
from mcp.server.auth.json_response import PydanticJSONResponse
from mcp.server.auth.provider import (
AuthorizationErrorCode,
AuthorizationParams,
AuthorizeError,
OAuthAuthorizationServerProvider,
construct_redirect_uri,
)
from mcp.shared.auth import InvalidRedirectUriError, InvalidScopeError
logger = logging.getLogger(__name__)
class AuthorizationRequest(BaseModel):
# See https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.1
client_id: str = Field(..., description="The client ID")
redirect_uri: AnyUrl | None = Field(None, description="URL to redirect to after authorization")
# see OAuthClientMetadata; we only support `code`
response_type: Literal["code"] = Field(..., description="Must be 'code' for authorization code flow")
code_challenge: str = Field(..., description="PKCE code challenge")
code_challenge_method: Literal["S256"] = Field("S256", description="PKCE code challenge method, must be S256")
state: str | None = Field(None, description="Optional state parameter")
scope: str | None = Field(
None,
description="Optional scope; if specified, should be a space-separated list of scope strings",
)
resource: str | None = Field(
None,
description="RFC 8707 resource indicator - the MCP server this token will be used with",
)
class AuthorizationErrorResponse(BaseModel):
error: AuthorizationErrorCode
error_description: str | None
error_uri: AnyUrl | None = None
# must be set if provided in the request
state: str | None = None
def best_effort_extract_string(key: str, params: None | FormData | QueryParams) -> str | None:
if params is None: # pragma: no cover
return None
value = params.get(key)
if isinstance(value, str):
return value
return None
class AnyUrlModel(RootModel[AnyUrl]):
root: AnyUrl
@dataclass
class AuthorizationHandler:
provider: OAuthAuthorizationServerProvider[Any, Any, Any]
async def handle(self, request: Request) -> Response:
# implements authorization requests for grant_type=code;
# see https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.1
state = None
redirect_uri = None
client = None
params = None
async def error_response(
error: AuthorizationErrorCode,
error_description: str | None,
attempt_load_client: bool = True,
):
# Error responses take two different formats:
# 1. The request has a valid client ID & redirect_uri: we issue a redirect
# back to the redirect_uri with the error response fields as query
# parameters. This allows the client to be notified of the error.
# 2. Otherwise, we return an error response directly to the end user;
# we choose to do so in JSON, but this is left undefined in the
# specification.
# See https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.2.1
#
# This logic is a bit awkward to handle, because the error might be thrown
# very early in request validation, before we've done the usual Pydantic
# validation, loaded the client, etc. To handle this, error_response()
# contains fallback logic which attempts to load the parameters directly
# from the request.
nonlocal client, redirect_uri, state
if client is None and attempt_load_client:
# make last-ditch attempt to load the client
client_id = best_effort_extract_string("client_id", params)
client = await self.provider.get_client(client_id) if client_id else None
if redirect_uri is None and client:
# make last-ditch effort to load the redirect uri
try:
if params is not None and "redirect_uri" not in params:
raw_redirect_uri = None
else:
raw_redirect_uri = AnyUrlModel.model_validate(
best_effort_extract_string("redirect_uri", params)
).root
redirect_uri = client.validate_redirect_uri(raw_redirect_uri)
except (ValidationError, InvalidRedirectUriError):
# if the redirect URI is invalid, ignore it & just return the
# initial error
pass
# the error response MUST contain the state specified by the client, if any
if state is None: # pragma: no cover
# make last-ditch effort to load state
state = best_effort_extract_string("state", params)
error_resp = AuthorizationErrorResponse(
error=error,
error_description=error_description,
state=state,
)
if redirect_uri and client:
return RedirectResponse(
url=construct_redirect_uri(str(redirect_uri), **error_resp.model_dump(exclude_none=True)),
status_code=302,
headers={"Cache-Control": "no-store"},
)
else:
return PydanticJSONResponse(
status_code=400,
content=error_resp,
headers={"Cache-Control": "no-store"},
)
try:
# Parse request parameters
if request.method == "GET":
# Convert query_params to dict for pydantic validation
params = request.query_params
else:
# Parse form data for POST requests
params = await request.form()
# Save state if it exists, even before validation
state = best_effort_extract_string("state", params)
try:
auth_request = AuthorizationRequest.model_validate(params)
state = auth_request.state # Update with validated state
except ValidationError as validation_error:
error: AuthorizationErrorCode = "invalid_request"
for e in validation_error.errors():
if e["loc"] == ("response_type",) and e["type"] == "literal_error":
error = "unsupported_response_type"
break
return await error_response(error, stringify_pydantic_error(validation_error))
# Get client information
client = await self.provider.get_client(
auth_request.client_id,
)
if not client:
# For client_id validation errors, return direct error (no redirect)
return await error_response(
error="invalid_request",
error_description=f"Client ID '{auth_request.client_id}' not found",
attempt_load_client=False,
)
# Validate redirect_uri against client's registered URIs
try:
redirect_uri = client.validate_redirect_uri(auth_request.redirect_uri)
except InvalidRedirectUriError as validation_error:
# For redirect_uri validation errors, return direct error (no redirect)
return await error_response(
error="invalid_request",
error_description=validation_error.message,
)
# Validate scope - for scope errors, we can redirect
try:
scopes = client.validate_scope(auth_request.scope)
except InvalidScopeError as validation_error:
# For scope errors, redirect with error parameters
return await error_response(
error="invalid_scope",
error_description=validation_error.message,
)
# Setup authorization parameters
auth_params = AuthorizationParams(
state=state,
scopes=scopes,
code_challenge=auth_request.code_challenge,
redirect_uri=redirect_uri,
redirect_uri_provided_explicitly=auth_request.redirect_uri is not None,
resource=auth_request.resource, # RFC 8707
)
try:
# Let the provider pick the next URI to redirect to
return RedirectResponse(
url=await self.provider.authorize(
client,
auth_params,
),
status_code=302,
headers={"Cache-Control": "no-store"},
)
except AuthorizeError as e:
# Handle authorization errors as defined in RFC 6749 Section 4.1.2.1
return await error_response(error=e.error, error_description=e.error_description)
except Exception as validation_error: # pragma: no cover
# Catch-all for unexpected errors
logger.exception("Unexpected error in authorization_handler", exc_info=validation_error)
return await error_response(error="server_error", error_description="An unexpected error occurred")
@@ -0,0 +1,29 @@
from dataclasses import dataclass
from starlette.requests import Request
from starlette.responses import Response
from mcp.server.auth.json_response import PydanticJSONResponse
from mcp.shared.auth import OAuthMetadata, ProtectedResourceMetadata
@dataclass
class MetadataHandler:
metadata: OAuthMetadata
async def handle(self, request: Request) -> Response:
return PydanticJSONResponse(
content=self.metadata,
headers={"Cache-Control": "public, max-age=3600"}, # Cache for 1 hour
)
@dataclass
class ProtectedResourceMetadataHandler:
metadata: ProtectedResourceMetadata
async def handle(self, request: Request) -> Response:
return PydanticJSONResponse(
content=self.metadata,
headers={"Cache-Control": "public, max-age=3600"}, # Cache for 1 hour
)
@@ -0,0 +1,136 @@
import secrets
import time
from dataclasses import dataclass
from typing import Any
from uuid import uuid4
from pydantic import BaseModel, RootModel, ValidationError
from starlette.requests import Request
from starlette.responses import Response
from mcp.server.auth.errors import stringify_pydantic_error
from mcp.server.auth.json_response import PydanticJSONResponse
from mcp.server.auth.provider import OAuthAuthorizationServerProvider, RegistrationError, RegistrationErrorCode
from mcp.server.auth.settings import ClientRegistrationOptions
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
class RegistrationRequest(RootModel[OAuthClientMetadata]):
# this wrapper is a no-op; it's just to separate out the types exposed to the
# provider from what we use in the HTTP handler
root: OAuthClientMetadata
class RegistrationErrorResponse(BaseModel):
error: RegistrationErrorCode
error_description: str | None
@dataclass
class RegistrationHandler:
provider: OAuthAuthorizationServerProvider[Any, Any, Any]
options: ClientRegistrationOptions
async def handle(self, request: Request) -> Response:
# Implements dynamic client registration as defined in https://datatracker.ietf.org/doc/html/rfc7591#section-3.1
try:
# Parse request body as JSON
body = await request.json()
client_metadata = OAuthClientMetadata.model_validate(body)
# Scope validation is handled below
except ValidationError as validation_error:
return PydanticJSONResponse(
content=RegistrationErrorResponse(
error="invalid_client_metadata",
error_description=stringify_pydantic_error(validation_error),
),
status_code=400,
)
client_id = str(uuid4())
# If auth method is None, default to client_secret_post
if client_metadata.token_endpoint_auth_method is None:
client_metadata.token_endpoint_auth_method = "client_secret_post"
client_secret = None
if client_metadata.token_endpoint_auth_method != "none": # pragma: no branch
# cryptographically secure random 32-byte hex string
client_secret = secrets.token_hex(32)
if client_metadata.scope is None and self.options.default_scopes is not None:
client_metadata.scope = " ".join(self.options.default_scopes)
elif client_metadata.scope is not None and self.options.valid_scopes is not None:
requested_scopes = set(client_metadata.scope.split())
valid_scopes = set(self.options.valid_scopes)
if not requested_scopes.issubset(valid_scopes): # pragma: no branch
return PydanticJSONResponse(
content=RegistrationErrorResponse(
error="invalid_client_metadata",
error_description="Requested scopes are not valid: "
f"{', '.join(requested_scopes - valid_scopes)}",
),
status_code=400,
)
if not {"authorization_code", "refresh_token"}.issubset(set(client_metadata.grant_types)):
return PydanticJSONResponse(
content=RegistrationErrorResponse(
error="invalid_client_metadata",
error_description="grant_types must be authorization_code and refresh_token",
),
status_code=400,
)
# The MCP spec requires servers to use the authorization `code` flow
# with PKCE
if "code" not in client_metadata.response_types:
return PydanticJSONResponse(
content=RegistrationErrorResponse(
error="invalid_client_metadata",
error_description="response_types must include 'code' for authorization_code grant",
),
status_code=400,
)
client_id_issued_at = int(time.time())
client_secret_expires_at = (
client_id_issued_at + self.options.client_secret_expiry_seconds
if self.options.client_secret_expiry_seconds is not None
else None
)
client_info = OAuthClientInformationFull(
client_id=client_id,
client_id_issued_at=client_id_issued_at,
client_secret=client_secret,
client_secret_expires_at=client_secret_expires_at,
# passthrough information from the client request
redirect_uris=client_metadata.redirect_uris,
token_endpoint_auth_method=client_metadata.token_endpoint_auth_method,
grant_types=client_metadata.grant_types,
response_types=client_metadata.response_types,
client_name=client_metadata.client_name,
client_uri=client_metadata.client_uri,
logo_uri=client_metadata.logo_uri,
scope=client_metadata.scope,
contacts=client_metadata.contacts,
tos_uri=client_metadata.tos_uri,
policy_uri=client_metadata.policy_uri,
jwks_uri=client_metadata.jwks_uri,
jwks=client_metadata.jwks,
software_id=client_metadata.software_id,
software_version=client_metadata.software_version,
)
try:
# Register client
await self.provider.register_client(client_info)
# Return client information
return PydanticJSONResponse(content=client_info, status_code=201)
except RegistrationError as e:
# Handle registration errors as defined in RFC 7591 Section 3.2.2
return PydanticJSONResponse(
content=RegistrationErrorResponse(error=e.error, error_description=e.error_description),
status_code=400,
)
@@ -0,0 +1,91 @@
from dataclasses import dataclass
from functools import partial
from typing import Any, Literal
from pydantic import BaseModel, ValidationError
from starlette.requests import Request
from starlette.responses import Response
from mcp.server.auth.errors import (
stringify_pydantic_error,
)
from mcp.server.auth.json_response import PydanticJSONResponse
from mcp.server.auth.middleware.client_auth import AuthenticationError, ClientAuthenticator
from mcp.server.auth.provider import AccessToken, OAuthAuthorizationServerProvider, RefreshToken
class RevocationRequest(BaseModel):
"""
# See https://datatracker.ietf.org/doc/html/rfc7009#section-2.1
"""
token: str
token_type_hint: Literal["access_token", "refresh_token"] | None = None
client_id: str
client_secret: str | None
class RevocationErrorResponse(BaseModel):
error: Literal["invalid_request", "unauthorized_client"]
error_description: str | None = None
@dataclass
class RevocationHandler:
provider: OAuthAuthorizationServerProvider[Any, Any, Any]
client_authenticator: ClientAuthenticator
async def handle(self, request: Request) -> Response:
"""
Handler for the OAuth 2.0 Token Revocation endpoint.
"""
try:
client = await self.client_authenticator.authenticate_request(request)
except AuthenticationError as e: # pragma: no cover
return PydanticJSONResponse(
status_code=401,
content=RevocationErrorResponse(
error="unauthorized_client",
error_description=e.message,
),
)
try:
form_data = await request.form()
revocation_request = RevocationRequest.model_validate(dict(form_data))
except ValidationError as e:
return PydanticJSONResponse(
status_code=400,
content=RevocationErrorResponse(
error="invalid_request",
error_description=stringify_pydantic_error(e),
),
)
loaders = [
self.provider.load_access_token,
partial(self.provider.load_refresh_token, client),
]
if revocation_request.token_type_hint == "refresh_token": # pragma: no cover
loaders = reversed(loaders)
token: None | AccessToken | RefreshToken = None
for loader in loaders:
token = await loader(revocation_request.token)
if token is not None:
break
# if token is not found, just return HTTP 200 per the RFC
if token and token.client_id == client.client_id:
# Revoke token; provider is not meant to be able to do validation
# at this point that would result in an error
await self.provider.revoke_token(token)
# Return successful empty response
return Response(
status_code=200,
headers={
"Cache-Control": "no-store",
"Pragma": "no-cache",
},
)
@@ -0,0 +1,241 @@
import base64
import hashlib
import time
from dataclasses import dataclass
from typing import Annotated, Any, Literal
from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, RootModel, ValidationError
from starlette.requests import Request
from mcp.server.auth.errors import stringify_pydantic_error
from mcp.server.auth.json_response import PydanticJSONResponse
from mcp.server.auth.middleware.client_auth import AuthenticationError, ClientAuthenticator
from mcp.server.auth.provider import OAuthAuthorizationServerProvider, TokenError, TokenErrorCode
from mcp.shared.auth import OAuthToken
class AuthorizationCodeRequest(BaseModel):
# See https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.3
grant_type: Literal["authorization_code"]
code: str = Field(..., description="The authorization code")
redirect_uri: AnyUrl | None = Field(None, description="Must be the same as redirect URI provided in /authorize")
client_id: str
# we use the client_secret param, per https://datatracker.ietf.org/doc/html/rfc6749#section-2.3.1
client_secret: str | None = None
# See https://datatracker.ietf.org/doc/html/rfc7636#section-4.5
code_verifier: str = Field(..., description="PKCE code verifier")
# RFC 8707 resource indicator
resource: str | None = Field(None, description="Resource indicator for the token")
class RefreshTokenRequest(BaseModel):
# See https://datatracker.ietf.org/doc/html/rfc6749#section-6
grant_type: Literal["refresh_token"]
refresh_token: str = Field(..., description="The refresh token")
scope: str | None = Field(None, description="Optional scope parameter")
client_id: str
# we use the client_secret param, per https://datatracker.ietf.org/doc/html/rfc6749#section-2.3.1
client_secret: str | None = None
# RFC 8707 resource indicator
resource: str | None = Field(None, description="Resource indicator for the token")
class TokenRequest(
RootModel[
Annotated[
AuthorizationCodeRequest | RefreshTokenRequest,
Field(discriminator="grant_type"),
]
]
):
root: Annotated[
AuthorizationCodeRequest | RefreshTokenRequest,
Field(discriminator="grant_type"),
]
class TokenErrorResponse(BaseModel):
"""
See https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
"""
error: TokenErrorCode
error_description: str | None = None
error_uri: AnyHttpUrl | None = None
class TokenSuccessResponse(RootModel[OAuthToken]):
# this is just a wrapper over OAuthToken; the only reason we do this
# is to have some separation between the HTTP response type, and the
# type returned by the provider
root: OAuthToken
@dataclass
class TokenHandler:
provider: OAuthAuthorizationServerProvider[Any, Any, Any]
client_authenticator: ClientAuthenticator
def response(self, obj: TokenSuccessResponse | TokenErrorResponse):
status_code = 200
if isinstance(obj, TokenErrorResponse):
status_code = 400
return PydanticJSONResponse(
content=obj,
status_code=status_code,
headers={
"Cache-Control": "no-store",
"Pragma": "no-cache",
},
)
async def handle(self, request: Request):
try:
client_info = await self.client_authenticator.authenticate_request(request)
except AuthenticationError as e:
# Authentication failures should return 401
return PydanticJSONResponse(
content=TokenErrorResponse(
error="unauthorized_client",
error_description=e.message,
),
status_code=401,
headers={
"Cache-Control": "no-store",
"Pragma": "no-cache",
},
)
try:
form_data = await request.form()
token_request = TokenRequest.model_validate(dict(form_data)).root
except ValidationError as validation_error: # pragma: no cover
return self.response(
TokenErrorResponse(
error="invalid_request",
error_description=stringify_pydantic_error(validation_error),
)
)
if token_request.grant_type not in client_info.grant_types: # pragma: no cover
return self.response(
TokenErrorResponse(
error="unsupported_grant_type",
error_description=(f"Unsupported grant type (supported grant types are {client_info.grant_types})"),
)
)
tokens: OAuthToken
match token_request:
case AuthorizationCodeRequest():
auth_code = await self.provider.load_authorization_code(client_info, token_request.code)
if auth_code is None or auth_code.client_id != token_request.client_id:
# if code belongs to different client, pretend it doesn't exist
return self.response(
TokenErrorResponse(
error="invalid_grant",
error_description="authorization code does not exist",
)
)
# make auth codes expire after a deadline
# see https://datatracker.ietf.org/doc/html/rfc6749#section-10.5
if auth_code.expires_at < time.time():
return self.response(
TokenErrorResponse(
error="invalid_grant",
error_description="authorization code has expired",
)
)
# verify redirect_uri doesn't change between /authorize and /tokens
# see https://datatracker.ietf.org/doc/html/rfc6749#section-10.6
if auth_code.redirect_uri_provided_explicitly:
authorize_request_redirect_uri = auth_code.redirect_uri
else: # pragma: no cover
authorize_request_redirect_uri = None
# Convert both sides to strings for comparison to handle AnyUrl vs string issues
token_redirect_str = str(token_request.redirect_uri) if token_request.redirect_uri is not None else None
auth_redirect_str = (
str(authorize_request_redirect_uri) if authorize_request_redirect_uri is not None else None
)
if token_redirect_str != auth_redirect_str:
return self.response(
TokenErrorResponse(
error="invalid_request",
error_description=("redirect_uri did not match the one used when creating auth code"),
)
)
# Verify PKCE code verifier
sha256 = hashlib.sha256(token_request.code_verifier.encode()).digest()
hashed_code_verifier = base64.urlsafe_b64encode(sha256).decode().rstrip("=")
if hashed_code_verifier != auth_code.code_challenge:
# see https://datatracker.ietf.org/doc/html/rfc7636#section-4.6
return self.response(
TokenErrorResponse(
error="invalid_grant",
error_description="incorrect code_verifier",
)
)
try:
# Exchange authorization code for tokens
tokens = await self.provider.exchange_authorization_code(client_info, auth_code)
except TokenError as e:
return self.response(
TokenErrorResponse(
error=e.error,
error_description=e.error_description,
)
)
case RefreshTokenRequest(): # pragma: no cover
refresh_token = await self.provider.load_refresh_token(client_info, token_request.refresh_token)
if refresh_token is None or refresh_token.client_id != token_request.client_id:
# if token belongs to different client, pretend it doesn't exist
return self.response(
TokenErrorResponse(
error="invalid_grant",
error_description="refresh token does not exist",
)
)
if refresh_token.expires_at and refresh_token.expires_at < time.time():
# if the refresh token has expired, pretend it doesn't exist
return self.response(
TokenErrorResponse(
error="invalid_grant",
error_description="refresh token has expired",
)
)
# Parse scopes if provided
scopes = token_request.scope.split(" ") if token_request.scope else refresh_token.scopes
for scope in scopes:
if scope not in refresh_token.scopes:
return self.response(
TokenErrorResponse(
error="invalid_scope",
error_description=(f"cannot request scope `{scope}` not provided by refresh token"),
)
)
try:
# Exchange refresh token for new tokens
tokens = await self.provider.exchange_refresh_token(client_info, refresh_token, scopes)
except TokenError as e:
return self.response(
TokenErrorResponse(
error=e.error,
error_description=e.error_description,
)
)
return self.response(TokenSuccessResponse(root=tokens))
@@ -0,0 +1,10 @@
from typing import Any
from starlette.responses import JSONResponse
class PydanticJSONResponse(JSONResponse):
# use pydantic json serialization instead of the stock `json.dumps`,
# so that we can handle serializing pydantic models like AnyHttpUrl
def render(self, content: Any) -> bytes:
return content.model_dump_json(exclude_none=True).encode("utf-8")
@@ -0,0 +1,3 @@
"""
Middleware for MCP authorization.
"""
@@ -0,0 +1,48 @@
import contextvars
from starlette.types import ASGIApp, Receive, Scope, Send
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
from mcp.server.auth.provider import AccessToken
# Create a contextvar to store the authenticated user
# The default is None, indicating no authenticated user is present
auth_context_var = contextvars.ContextVar[AuthenticatedUser | None]("auth_context", default=None)
def get_access_token() -> AccessToken | None:
"""
Get the access token from the current context.
Returns:
The access token if an authenticated user is available, None otherwise.
"""
auth_user = auth_context_var.get()
return auth_user.access_token if auth_user else None
class AuthContextMiddleware:
"""
Middleware that extracts the authenticated user from the request
and sets it in a contextvar for easy access throughout the request lifecycle.
This middleware should be added after the AuthenticationMiddleware in the
middleware stack to ensure that the user is properly authenticated before
being stored in the context.
"""
def __init__(self, app: ASGIApp):
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send):
user = scope.get("user")
if isinstance(user, AuthenticatedUser):
# Set the authenticated user in the contextvar
token = auth_context_var.set(user)
try:
await self.app(scope, receive, send)
finally:
auth_context_var.reset(token)
else:
# No authenticated user, just process the request
await self.app(scope, receive, send)
@@ -0,0 +1,128 @@
import json
import time
from typing import Any
from pydantic import AnyHttpUrl
from starlette.authentication import AuthCredentials, AuthenticationBackend, SimpleUser
from starlette.requests import HTTPConnection
from starlette.types import Receive, Scope, Send
from mcp.server.auth.provider import AccessToken, TokenVerifier
class AuthenticatedUser(SimpleUser):
"""User with authentication info."""
def __init__(self, auth_info: AccessToken):
super().__init__(auth_info.client_id)
self.access_token = auth_info
self.scopes = auth_info.scopes
class BearerAuthBackend(AuthenticationBackend):
"""
Authentication backend that validates Bearer tokens using a TokenVerifier.
"""
def __init__(self, token_verifier: TokenVerifier):
self.token_verifier = token_verifier
async def authenticate(self, conn: HTTPConnection):
auth_header = next(
(conn.headers.get(key) for key in conn.headers if key.lower() == "authorization"),
None,
)
if not auth_header or not auth_header.lower().startswith("bearer "):
return None
token = auth_header[7:] # Remove "Bearer " prefix
# Validate the token with the verifier
auth_info = await self.token_verifier.verify_token(token)
if not auth_info:
return None
if auth_info.expires_at and auth_info.expires_at < int(time.time()):
return None
return AuthCredentials(auth_info.scopes), AuthenticatedUser(auth_info)
class RequireAuthMiddleware:
"""
Middleware that requires a valid Bearer token in the Authorization header.
This will validate the token with the auth provider and store the resulting
auth info in the request state.
"""
def __init__(
self,
app: Any,
required_scopes: list[str],
resource_metadata_url: AnyHttpUrl | None = None,
):
"""
Initialize the middleware.
Args:
app: ASGI application
required_scopes: List of scopes that the token must have
resource_metadata_url: Optional protected resource metadata URL for WWW-Authenticate header
"""
self.app = app
self.required_scopes = required_scopes
self.resource_metadata_url = resource_metadata_url
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
auth_user = scope.get("user")
if not isinstance(auth_user, AuthenticatedUser):
await self._send_auth_error(
send, status_code=401, error="invalid_token", description="Authentication required"
)
return
auth_credentials = scope.get("auth")
for required_scope in self.required_scopes:
# auth_credentials should always be provided; this is just paranoia
if auth_credentials is None or required_scope not in auth_credentials.scopes:
await self._send_auth_error(
send, status_code=403, error="insufficient_scope", description=f"Required scope: {required_scope}"
)
return
await self.app(scope, receive, send)
async def _send_auth_error(self, send: Send, status_code: int, error: str, description: str) -> None:
"""Send an authentication error response with WWW-Authenticate header."""
# Build WWW-Authenticate header value
www_auth_parts = [f'error="{error}"', f'error_description="{description}"']
if self.resource_metadata_url: # pragma: no cover
www_auth_parts.append(f'resource_metadata="{self.resource_metadata_url}"')
www_authenticate = f"Bearer {', '.join(www_auth_parts)}"
# Send response
body = {"error": error, "error_description": description}
body_bytes = json.dumps(body).encode()
await send(
{
"type": "http.response.start",
"status": status_code,
"headers": [
(b"content-type", b"application/json"),
(b"content-length", str(len(body_bytes)).encode()),
(b"www-authenticate", www_authenticate.encode()),
],
}
)
await send(
{
"type": "http.response.body",
"body": body_bytes,
}
)
@@ -0,0 +1,115 @@
import base64
import binascii
import hmac
import time
from typing import Any
from urllib.parse import unquote
from starlette.requests import Request
from mcp.server.auth.provider import OAuthAuthorizationServerProvider
from mcp.shared.auth import OAuthClientInformationFull
class AuthenticationError(Exception):
def __init__(self, message: str):
self.message = message # pragma: no cover
class ClientAuthenticator:
"""
ClientAuthenticator is a callable which validates requests from a client
application, used to verify /token calls.
If, during registration, the client requested to be issued a secret, the
authenticator asserts that /token calls must be authenticated with
that same token.
NOTE: clients can opt for no authentication during registration, in which case this
logic is skipped.
"""
def __init__(self, provider: OAuthAuthorizationServerProvider[Any, Any, Any]):
"""
Initialize the dependency.
Args:
provider: Provider to look up client information
"""
self.provider = provider
async def authenticate_request(self, request: Request) -> OAuthClientInformationFull:
"""
Authenticate a client from an HTTP request.
Extracts client credentials from the appropriate location based on the
client's registered authentication method and validates them.
Args:
request: The HTTP request containing client credentials
Returns:
The authenticated client information
Raises:
AuthenticationError: If authentication fails
"""
form_data = await request.form()
client_id = form_data.get("client_id")
if not client_id:
raise AuthenticationError("Missing client_id")
client = await self.provider.get_client(str(client_id))
if not client:
raise AuthenticationError("Invalid client_id") # pragma: no cover
request_client_secret: str | None = None
auth_header = request.headers.get("Authorization", "")
if client.token_endpoint_auth_method == "client_secret_basic":
if not auth_header.startswith("Basic "):
raise AuthenticationError("Missing or invalid Basic authentication in Authorization header")
try:
encoded_credentials = auth_header[6:] # Remove "Basic " prefix
decoded = base64.b64decode(encoded_credentials).decode("utf-8")
if ":" not in decoded:
raise ValueError("Invalid Basic auth format")
basic_client_id, request_client_secret = decoded.split(":", 1)
# URL-decode both parts per RFC 6749 Section 2.3.1
basic_client_id = unquote(basic_client_id)
request_client_secret = unquote(request_client_secret)
if basic_client_id != client_id:
raise AuthenticationError("Client ID mismatch in Basic auth")
except (ValueError, UnicodeDecodeError, binascii.Error):
raise AuthenticationError("Invalid Basic authentication header")
elif client.token_endpoint_auth_method == "client_secret_post":
raw_form_data = form_data.get("client_secret")
# form_data.get() can return a UploadFile or None, so we need to check if it's a string
if isinstance(raw_form_data, str):
request_client_secret = str(raw_form_data)
elif client.token_endpoint_auth_method == "none":
request_client_secret = None
else:
raise AuthenticationError( # pragma: no cover
f"Unsupported auth method: {client.token_endpoint_auth_method}"
)
# If client from the store expects a secret, validate that the request provides
# that secret
if client.client_secret: # pragma: no branch
if not request_client_secret:
raise AuthenticationError("Client secret is required") # pragma: no cover
# hmac.compare_digest requires that both arguments are either bytes or a `str` containing
# only ASCII characters. Since we do not control `request_client_secret`, we encode both
# arguments to bytes.
if not hmac.compare_digest(client.client_secret.encode(), request_client_secret.encode()):
raise AuthenticationError("Invalid client_secret") # pragma: no cover
if client.client_secret_expires_at and client.client_secret_expires_at < int(time.time()):
raise AuthenticationError("Client secret has expired") # pragma: no cover
return client
@@ -0,0 +1,301 @@
from dataclasses import dataclass
from typing import Generic, Literal, Protocol, TypeVar
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
from pydantic import AnyUrl, BaseModel
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
class AuthorizationParams(BaseModel):
state: str | None
scopes: list[str] | None
code_challenge: str
redirect_uri: AnyUrl
redirect_uri_provided_explicitly: bool
resource: str | None = None # RFC 8707 resource indicator
class AuthorizationCode(BaseModel):
code: str
scopes: list[str]
expires_at: float
client_id: str
code_challenge: str
redirect_uri: AnyUrl
redirect_uri_provided_explicitly: bool
resource: str | None = None # RFC 8707 resource indicator
class RefreshToken(BaseModel):
token: str
client_id: str
scopes: list[str]
expires_at: int | None = None
class AccessToken(BaseModel):
token: str
client_id: str
scopes: list[str]
expires_at: int | None = None
resource: str | None = None # RFC 8707 resource indicator
RegistrationErrorCode = Literal[
"invalid_redirect_uri",
"invalid_client_metadata",
"invalid_software_statement",
"unapproved_software_statement",
]
@dataclass(frozen=True)
class RegistrationError(Exception):
error: RegistrationErrorCode
error_description: str | None = None
AuthorizationErrorCode = Literal[
"invalid_request",
"unauthorized_client",
"access_denied",
"unsupported_response_type",
"invalid_scope",
"server_error",
"temporarily_unavailable",
]
@dataclass(frozen=True)
class AuthorizeError(Exception):
error: AuthorizationErrorCode
error_description: str | None = None
TokenErrorCode = Literal[
"invalid_request",
"invalid_client",
"invalid_grant",
"unauthorized_client",
"unsupported_grant_type",
"invalid_scope",
]
@dataclass(frozen=True)
class TokenError(Exception):
error: TokenErrorCode
error_description: str | None = None
class TokenVerifier(Protocol):
"""Protocol for verifying bearer tokens."""
async def verify_token(self, token: str) -> AccessToken | None:
"""Verify a bearer token and return access info if valid."""
# NOTE: FastMCP doesn't render any of these types in the user response, so it's
# OK to add fields to subclasses which should not be exposed externally.
AuthorizationCodeT = TypeVar("AuthorizationCodeT", bound=AuthorizationCode)
RefreshTokenT = TypeVar("RefreshTokenT", bound=RefreshToken)
AccessTokenT = TypeVar("AccessTokenT", bound=AccessToken)
class OAuthAuthorizationServerProvider(Protocol, Generic[AuthorizationCodeT, RefreshTokenT, AccessTokenT]):
async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
"""
Retrieves client information by client ID.
Implementors MAY raise NotImplementedError if dynamic client registration is
disabled in ClientRegistrationOptions.
Args:
client_id: The ID of the client to retrieve.
Returns:
The client information, or None if the client does not exist.
"""
async def register_client(self, client_info: OAuthClientInformationFull) -> None:
"""
Saves client information as part of registering it.
Implementors MAY raise NotImplementedError if dynamic client registration is
disabled in ClientRegistrationOptions.
Args:
client_info: The client metadata to register.
Raises:
RegistrationError: If the client metadata is invalid.
"""
async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
"""
Called as part of the /authorize endpoint, and returns a URL that the client
will be redirected to.
Many MCP implementations will redirect to a third-party provider to perform
a second OAuth exchange with that provider. In this sort of setup, the client
has an OAuth connection with the MCP server, and the MCP server has an OAuth
connection with the 3rd-party provider. At the end of this flow, the client
should be redirected to the redirect_uri from params.redirect_uri.
+--------+ +------------+ +-------------------+
| | | | | |
| Client | --> | MCP Server | --> | 3rd Party OAuth |
| | | | | Server |
+--------+ +------------+ +-------------------+
| ^ |
+------------+ | | |
| | | | Redirect |
|redirect_uri|<-----+ +------------------+
| |
+------------+
Implementations will need to define another handler on the MCP server return
flow to perform the second redirect, and generate and store an authorization
code as part of completing the OAuth authorization step.
Implementations SHOULD generate an authorization code with at least 160 bits of
entropy,
and MUST generate an authorization code with at least 128 bits of entropy.
See https://datatracker.ietf.org/doc/html/rfc6749#section-10.10.
Args:
client: The client requesting authorization.
params: The parameters of the authorization request.
Returns:
A URL to redirect the client to for authorization.
Raises:
AuthorizeError: If the authorization request is invalid.
"""
...
async def load_authorization_code(
self, client: OAuthClientInformationFull, authorization_code: str
) -> AuthorizationCodeT | None:
"""
Loads an AuthorizationCode by its code.
Args:
client: The client that requested the authorization code.
authorization_code: The authorization code to get the challenge for.
Returns:
The AuthorizationCode, or None if not found
"""
...
async def exchange_authorization_code(
self, client: OAuthClientInformationFull, authorization_code: AuthorizationCodeT
) -> OAuthToken:
"""
Exchanges an authorization code for an access token and refresh token.
Args:
client: The client exchanging the authorization code.
authorization_code: The authorization code to exchange.
Returns:
The OAuth token, containing access and refresh tokens.
Raises:
TokenError: If the request is invalid
"""
...
async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshTokenT | None:
"""
Loads a RefreshToken by its token string.
Args:
client: The client that is requesting to load the refresh token.
refresh_token: The refresh token string to load.
Returns:
The RefreshToken object if found, or None if not found.
"""
...
async def exchange_refresh_token(
self,
client: OAuthClientInformationFull,
refresh_token: RefreshTokenT,
scopes: list[str],
) -> OAuthToken:
"""
Exchanges a refresh token for an access token and refresh token.
Implementations SHOULD rotate both the access token and refresh token.
Args:
client: The client exchanging the refresh token.
refresh_token: The refresh token to exchange.
scopes: Optional scopes to request with the new access token.
Returns:
The OAuth token, containing access and refresh tokens.
Raises:
TokenError: If the request is invalid
"""
...
async def load_access_token(self, token: str) -> AccessTokenT | None:
"""
Loads an access token by its token.
Args:
token: The access token to verify.
Returns:
The AuthInfo, or None if the token is invalid.
"""
async def revoke_token(
self,
token: AccessTokenT | RefreshTokenT,
) -> None:
"""
Revokes an access or refresh token.
If the given token is invalid or already revoked, this method should do nothing.
Implementations SHOULD revoke both the access token and its corresponding
refresh token, regardless of which of the access token or refresh token is
provided.
Args:
token: the token to revoke
"""
def construct_redirect_uri(redirect_uri_base: str, **params: str | None) -> str:
parsed_uri = urlparse(redirect_uri_base)
query_params = [(k, v) for k, vs in parse_qs(parsed_uri.query).items() for v in vs]
for k, v in params.items():
if v is not None:
query_params.append((k, v))
redirect_uri = urlunparse(parsed_uri._replace(query=urlencode(query_params)))
return redirect_uri
class ProviderTokenVerifier(TokenVerifier):
"""Token verifier that uses an OAuthAuthorizationServerProvider.
This is provided for backwards compatibility with existing auth_server_provider
configurations. For new implementations using AS/RS separation, consider using
the TokenVerifier protocol with a dedicated implementation like IntrospectionTokenVerifier.
"""
def __init__(self, provider: "OAuthAuthorizationServerProvider[AuthorizationCode, RefreshToken, AccessToken]"):
self.provider = provider
async def verify_token(self, token: str) -> AccessToken | None:
"""Verify token using the provider's load_access_token method."""
return await self.provider.load_access_token(token)
@@ -0,0 +1,253 @@
from collections.abc import Awaitable, Callable
from typing import Any
from urllib.parse import urlparse
from pydantic import AnyHttpUrl
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Route, request_response # type: ignore
from starlette.types import ASGIApp
from mcp.server.auth.handlers.authorize import AuthorizationHandler
from mcp.server.auth.handlers.metadata import MetadataHandler
from mcp.server.auth.handlers.register import RegistrationHandler
from mcp.server.auth.handlers.revoke import RevocationHandler
from mcp.server.auth.handlers.token import TokenHandler
from mcp.server.auth.middleware.client_auth import ClientAuthenticator
from mcp.server.auth.provider import OAuthAuthorizationServerProvider
from mcp.server.auth.settings import ClientRegistrationOptions, RevocationOptions
from mcp.server.streamable_http import MCP_PROTOCOL_VERSION_HEADER
from mcp.shared.auth import OAuthMetadata
def validate_issuer_url(url: AnyHttpUrl):
"""
Validate that the issuer URL meets OAuth 2.0 requirements.
Args:
url: The issuer URL to validate
Raises:
ValueError: If the issuer URL is invalid
"""
# RFC 8414 requires HTTPS, but we allow localhost HTTP for testing
if (
url.scheme != "https"
and url.host != "localhost"
and (url.host is not None and not url.host.startswith("127.0.0.1"))
):
raise ValueError("Issuer URL must be HTTPS") # pragma: no cover
# No fragments or query parameters allowed
if url.fragment:
raise ValueError("Issuer URL must not have a fragment") # pragma: no cover
if url.query:
raise ValueError("Issuer URL must not have a query string") # pragma: no cover
AUTHORIZATION_PATH = "/authorize"
TOKEN_PATH = "/token"
REGISTRATION_PATH = "/register"
REVOCATION_PATH = "/revoke"
def cors_middleware(
handler: Callable[[Request], Response | Awaitable[Response]],
allow_methods: list[str],
) -> ASGIApp:
cors_app = CORSMiddleware(
app=request_response(handler),
allow_origins="*",
allow_methods=allow_methods,
allow_headers=[MCP_PROTOCOL_VERSION_HEADER],
)
return cors_app
def create_auth_routes(
provider: OAuthAuthorizationServerProvider[Any, Any, Any],
issuer_url: AnyHttpUrl,
service_documentation_url: AnyHttpUrl | None = None,
client_registration_options: ClientRegistrationOptions | None = None,
revocation_options: RevocationOptions | None = None,
) -> list[Route]:
validate_issuer_url(issuer_url)
client_registration_options = client_registration_options or ClientRegistrationOptions()
revocation_options = revocation_options or RevocationOptions()
metadata = build_metadata(
issuer_url,
service_documentation_url,
client_registration_options,
revocation_options,
)
client_authenticator = ClientAuthenticator(provider)
# Create routes
# Allow CORS requests for endpoints meant to be hit by the OAuth client
# (with the client secret). This is intended to support things like MCP Inspector,
# where the client runs in a web browser.
routes = [
Route(
"/.well-known/oauth-authorization-server",
endpoint=cors_middleware(
MetadataHandler(metadata).handle,
["GET", "OPTIONS"],
),
methods=["GET", "OPTIONS"],
),
Route(
AUTHORIZATION_PATH,
# do not allow CORS for authorization endpoint;
# clients should just redirect to this
endpoint=AuthorizationHandler(provider).handle,
methods=["GET", "POST"],
),
Route(
TOKEN_PATH,
endpoint=cors_middleware(
TokenHandler(provider, client_authenticator).handle,
["POST", "OPTIONS"],
),
methods=["POST", "OPTIONS"],
),
]
if client_registration_options.enabled: # pragma: no branch
registration_handler = RegistrationHandler(
provider,
options=client_registration_options,
)
routes.append(
Route(
REGISTRATION_PATH,
endpoint=cors_middleware(
registration_handler.handle,
["POST", "OPTIONS"],
),
methods=["POST", "OPTIONS"],
)
)
if revocation_options.enabled: # pragma: no branch
revocation_handler = RevocationHandler(provider, client_authenticator)
routes.append(
Route(
REVOCATION_PATH,
endpoint=cors_middleware(
revocation_handler.handle,
["POST", "OPTIONS"],
),
methods=["POST", "OPTIONS"],
)
)
return routes
def build_metadata(
issuer_url: AnyHttpUrl,
service_documentation_url: AnyHttpUrl | None,
client_registration_options: ClientRegistrationOptions,
revocation_options: RevocationOptions,
) -> OAuthMetadata:
authorization_url = AnyHttpUrl(str(issuer_url).rstrip("/") + AUTHORIZATION_PATH)
token_url = AnyHttpUrl(str(issuer_url).rstrip("/") + TOKEN_PATH)
# Create metadata
metadata = OAuthMetadata(
issuer=issuer_url,
authorization_endpoint=authorization_url,
token_endpoint=token_url,
scopes_supported=client_registration_options.valid_scopes,
response_types_supported=["code"],
response_modes_supported=None,
grant_types_supported=["authorization_code", "refresh_token"],
token_endpoint_auth_methods_supported=["client_secret_post", "client_secret_basic"],
token_endpoint_auth_signing_alg_values_supported=None,
service_documentation=service_documentation_url,
ui_locales_supported=None,
op_policy_uri=None,
op_tos_uri=None,
introspection_endpoint=None,
code_challenge_methods_supported=["S256"],
)
# Add registration endpoint if supported
if client_registration_options.enabled: # pragma: no branch
metadata.registration_endpoint = AnyHttpUrl(str(issuer_url).rstrip("/") + REGISTRATION_PATH)
# Add revocation endpoint if supported
if revocation_options.enabled: # pragma: no branch
metadata.revocation_endpoint = AnyHttpUrl(str(issuer_url).rstrip("/") + REVOCATION_PATH)
metadata.revocation_endpoint_auth_methods_supported = ["client_secret_post", "client_secret_basic"]
return metadata
def build_resource_metadata_url(resource_server_url: AnyHttpUrl) -> AnyHttpUrl:
"""
Build RFC 9728 compliant protected resource metadata URL.
Inserts /.well-known/oauth-protected-resource between host and resource path
as specified in RFC 9728 §3.1.
Args:
resource_server_url: The resource server URL (e.g., https://example.com/mcp)
Returns:
The metadata URL (e.g., https://example.com/.well-known/oauth-protected-resource/mcp)
"""
parsed = urlparse(str(resource_server_url))
# Handle trailing slash: if path is just "/", treat as empty
resource_path = parsed.path if parsed.path != "/" else ""
return AnyHttpUrl(f"{parsed.scheme}://{parsed.netloc}/.well-known/oauth-protected-resource{resource_path}")
def create_protected_resource_routes(
resource_url: AnyHttpUrl,
authorization_servers: list[AnyHttpUrl],
scopes_supported: list[str] | None = None,
resource_name: str | None = None,
resource_documentation: AnyHttpUrl | None = None,
) -> list[Route]:
"""
Create routes for OAuth 2.0 Protected Resource Metadata (RFC 9728).
Args:
resource_url: The URL of this resource server
authorization_servers: List of authorization servers that can issue tokens
scopes_supported: Optional list of scopes supported by this resource
Returns:
List of Starlette routes for protected resource metadata
"""
from mcp.server.auth.handlers.metadata import ProtectedResourceMetadataHandler
from mcp.shared.auth import ProtectedResourceMetadata
metadata = ProtectedResourceMetadata(
resource=resource_url,
authorization_servers=authorization_servers,
scopes_supported=scopes_supported,
resource_name=resource_name,
resource_documentation=resource_documentation,
# bearer_methods_supported defaults to ["header"] in the model
)
handler = ProtectedResourceMetadataHandler(metadata)
# RFC 9728 §3.1: Register route at /.well-known/oauth-protected-resource + resource path
metadata_url = build_resource_metadata_url(resource_url)
# Extract just the path part for route registration
parsed = urlparse(str(metadata_url))
well_known_path = parsed.path
return [
Route(
well_known_path,
endpoint=cors_middleware(handler.handle, ["GET", "OPTIONS"]),
methods=["GET", "OPTIONS"],
)
]
@@ -0,0 +1,30 @@
from pydantic import AnyHttpUrl, BaseModel, Field
class ClientRegistrationOptions(BaseModel):
enabled: bool = False
client_secret_expiry_seconds: int | None = None
valid_scopes: list[str] | None = None
default_scopes: list[str] | None = None
class RevocationOptions(BaseModel):
enabled: bool = False
class AuthSettings(BaseModel):
issuer_url: AnyHttpUrl = Field(
...,
description="OAuth authorization server URL that issues tokens for this resource server.",
)
service_documentation_url: AnyHttpUrl | None = None
client_registration_options: ClientRegistrationOptions | None = None
revocation_options: RevocationOptions | None = None
required_scopes: list[str] | None = None
# Resource Server settings (when operating as RS only)
resource_server_url: AnyHttpUrl | None = Field(
...,
description="The URL of the MCP server to be used as the resource identifier "
"and base route to look up OAuth Protected Resource Metadata.",
)
@@ -0,0 +1,190 @@
"""Elicitation utilities for MCP servers."""
from __future__ import annotations
import types
from collections.abc import Sequence
from typing import Generic, Literal, TypeVar, Union, get_args, get_origin
from pydantic import BaseModel
from mcp.server.session import ServerSession
from mcp.types import RequestId
ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)
class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]):
"""Result when user accepts the elicitation."""
action: Literal["accept"] = "accept"
data: ElicitSchemaModelT
class DeclinedElicitation(BaseModel):
"""Result when user declines the elicitation."""
action: Literal["decline"] = "decline"
class CancelledElicitation(BaseModel):
"""Result when user cancels the elicitation."""
action: Literal["cancel"] = "cancel"
ElicitationResult = AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation
class AcceptedUrlElicitation(BaseModel):
"""Result when user accepts a URL mode elicitation."""
action: Literal["accept"] = "accept"
UrlElicitationResult = AcceptedUrlElicitation | DeclinedElicitation | CancelledElicitation
# Primitive types allowed in elicitation schemas
_ELICITATION_PRIMITIVE_TYPES = (str, int, float, bool)
def _validate_elicitation_schema(schema: type[BaseModel]) -> None:
"""Validate that a Pydantic model only contains primitive field types."""
for field_name, field_info in schema.model_fields.items():
annotation = field_info.annotation
if annotation is None or annotation is types.NoneType: # pragma: no cover
continue
elif _is_primitive_field(annotation):
continue
elif _is_string_sequence(annotation):
continue
else:
raise TypeError(
f"Elicitation schema field '{field_name}' must be a primitive type "
f"{_ELICITATION_PRIMITIVE_TYPES}, a sequence of strings (list[str], etc.), "
f"or Optional of these types. Nested models and complex types are not allowed."
)
def _is_string_sequence(annotation: type) -> bool:
"""Check if annotation is a sequence of strings (list[str], Sequence[str], etc)."""
origin = get_origin(annotation)
# Check if it's a sequence-like type with str elements
if origin:
try:
if issubclass(origin, Sequence):
args = get_args(annotation)
# Should have single str type arg
return len(args) == 1 and args[0] is str
except TypeError: # pragma: no cover
# origin is not a class, so it can't be a subclass of Sequence
pass
return False
def _is_primitive_field(annotation: type) -> bool:
"""Check if a field is a primitive type allowed in elicitation schemas."""
# Handle basic primitive types
if annotation in _ELICITATION_PRIMITIVE_TYPES:
return True
# Handle Union types
origin = get_origin(annotation)
if origin is Union or origin is types.UnionType:
args = get_args(annotation)
# All args must be primitive types, None, or string sequences
return all(
arg is types.NoneType or arg in _ELICITATION_PRIMITIVE_TYPES or _is_string_sequence(arg) for arg in args
)
return False
async def elicit_with_validation(
session: ServerSession,
message: str,
schema: type[ElicitSchemaModelT],
related_request_id: RequestId | None = None,
) -> ElicitationResult[ElicitSchemaModelT]:
"""Elicit information from the client/user with schema validation (form mode).
This method can be used to interactively ask for additional information from the
client within a tool's execution. The client might display the message to the
user and collect a response according to the provided schema. Or in case a
client is an agent, it might decide how to handle the elicitation -- either by asking
the user or automatically generating a response.
For sensitive data like credentials or OAuth flows, use elicit_url() instead.
"""
# Validate that schema only contains primitive types and fail loudly if not
_validate_elicitation_schema(schema)
json_schema = schema.model_json_schema()
result = await session.elicit_form(
message=message,
requestedSchema=json_schema,
related_request_id=related_request_id,
)
if result.action == "accept" and result.content is not None:
# Validate and parse the content using the schema
validated_data = schema.model_validate(result.content)
return AcceptedElicitation(data=validated_data)
elif result.action == "decline":
return DeclinedElicitation()
elif result.action == "cancel": # pragma: no cover
return CancelledElicitation()
else: # pragma: no cover
# This should never happen, but handle it just in case
raise ValueError(f"Unexpected elicitation action: {result.action}")
async def elicit_url(
session: ServerSession,
message: str,
url: str,
elicitation_id: str,
related_request_id: RequestId | None = None,
) -> UrlElicitationResult:
"""Elicit information from the user via out-of-band URL navigation (URL mode).
This method directs the user to an external URL where sensitive interactions can
occur without passing data through the MCP client. Use this for:
- Collecting sensitive credentials (API keys, passwords)
- OAuth authorization flows with third-party services
- Payment and subscription flows
- Any interaction where data should not pass through the LLM context
The response indicates whether the user consented to navigate to the URL.
The actual interaction happens out-of-band. When the elicitation completes,
the server should send an ElicitCompleteNotification to notify the client.
Args:
session: The server session
message: Human-readable explanation of why the interaction is needed
url: The URL the user should navigate to
elicitation_id: Unique identifier for tracking this elicitation
related_request_id: Optional ID of the request that triggered this elicitation
Returns:
UrlElicitationResult indicating accept, decline, or cancel
"""
result = await session.elicit_url(
message=message,
url=url,
elicitation_id=elicitation_id,
related_request_id=related_request_id,
)
if result.action == "accept":
return AcceptedUrlElicitation()
elif result.action == "decline":
return DeclinedElicitation()
elif result.action == "cancel":
return CancelledElicitation()
else: # pragma: no cover
# This should never happen, but handle it just in case
raise ValueError(f"Unexpected elicitation action: {result.action}")
@@ -0,0 +1,11 @@
"""
Server-side experimental features.
WARNING: These APIs are experimental and may change without notice.
Import directly from submodules:
- mcp.server.experimental.task_context.ServerTaskContext
- mcp.server.experimental.task_support.TaskSupport
- mcp.server.experimental.task_result_handler.TaskResultHandler
- mcp.server.experimental.request_context.Experimental
"""
@@ -0,0 +1,238 @@
"""
Experimental request context features.
This module provides the Experimental class which gives access to experimental
features within a request context, such as task-augmented request handling.
WARNING: These APIs are experimental and may change without notice.
"""
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any
from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.experimental.task_support import TaskSupport
from mcp.server.session import ServerSession
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.helpers import MODEL_IMMEDIATE_RESPONSE_KEY, is_terminal
from mcp.types import (
METHOD_NOT_FOUND,
TASK_FORBIDDEN,
TASK_REQUIRED,
ClientCapabilities,
CreateTaskResult,
ErrorData,
Result,
TaskExecutionMode,
TaskMetadata,
Tool,
)
@dataclass
class Experimental:
"""
Experimental features context for task-augmented requests.
Provides helpers for validating task execution compatibility and
running tasks with automatic lifecycle management.
WARNING: This API is experimental and may change without notice.
"""
task_metadata: TaskMetadata | None = None
_client_capabilities: ClientCapabilities | None = field(default=None, repr=False)
_session: ServerSession | None = field(default=None, repr=False)
_task_support: TaskSupport | None = field(default=None, repr=False)
@property
def is_task(self) -> bool:
"""Check if this request is task-augmented."""
return self.task_metadata is not None
@property
def client_supports_tasks(self) -> bool:
"""Check if the client declared task support."""
if self._client_capabilities is None:
return False
return self._client_capabilities.tasks is not None
def validate_task_mode(
self,
tool_task_mode: TaskExecutionMode | None,
*,
raise_error: bool = True,
) -> ErrorData | None:
"""
Validate that the request is compatible with the tool's task execution mode.
Per MCP spec:
- "required": Clients MUST invoke as task. Server returns -32601 if not.
- "forbidden" (or None): Clients MUST NOT invoke as task. Server returns -32601 if they do.
- "optional": Either is acceptable.
Args:
tool_task_mode: The tool's execution.taskSupport value
("forbidden", "optional", "required", or None)
raise_error: If True, raises McpError on validation failure. If False, returns ErrorData.
Returns:
None if valid, ErrorData if invalid and raise_error=False
Raises:
McpError: If invalid and raise_error=True
"""
mode = tool_task_mode or TASK_FORBIDDEN
error: ErrorData | None = None
if mode == TASK_REQUIRED and not self.is_task:
error = ErrorData(
code=METHOD_NOT_FOUND,
message="This tool requires task-augmented invocation",
)
elif mode == TASK_FORBIDDEN and self.is_task:
error = ErrorData(
code=METHOD_NOT_FOUND,
message="This tool does not support task-augmented invocation",
)
if error is not None and raise_error:
raise McpError(error)
return error
def validate_for_tool(
self,
tool: Tool,
*,
raise_error: bool = True,
) -> ErrorData | None:
"""
Validate that the request is compatible with the given tool.
Convenience wrapper around validate_task_mode that extracts the mode from a Tool.
Args:
tool: The Tool definition
raise_error: If True, raises McpError on validation failure.
Returns:
None if valid, ErrorData if invalid and raise_error=False
"""
mode = tool.execution.taskSupport if tool.execution else None
return self.validate_task_mode(mode, raise_error=raise_error)
def can_use_tool(self, tool_task_mode: TaskExecutionMode | None) -> bool:
"""
Check if this client can use a tool with the given task mode.
Useful for filtering tool lists or providing warnings.
Returns False if tool requires "required" but client doesn't support tasks.
Args:
tool_task_mode: The tool's execution.taskSupport value
Returns:
True if the client can use this tool, False otherwise
"""
mode = tool_task_mode or TASK_FORBIDDEN
if mode == TASK_REQUIRED and not self.client_supports_tasks:
return False
return True
async def run_task(
self,
work: Callable[[ServerTaskContext], Awaitable[Result]],
*,
task_id: str | None = None,
model_immediate_response: str | None = None,
) -> CreateTaskResult:
"""
Create a task, spawn background work, and return CreateTaskResult immediately.
This is the recommended way to handle task-augmented tool calls. It:
1. Creates a task in the store
2. Spawns the work function in a background task
3. Returns CreateTaskResult immediately
The work function receives a ServerTaskContext with:
- elicit() for sending elicitation requests
- create_message() for sampling requests
- update_status() for progress updates
- complete()/fail() for finishing the task
When work() returns a Result, the task is auto-completed with that result.
If work() raises an exception, the task is auto-failed.
Args:
work: Async function that does the actual work
task_id: Optional task ID (generated if not provided)
model_immediate_response: Optional string to include in _meta as
io.modelcontextprotocol/model-immediate-response
Returns:
CreateTaskResult to return to the client
Raises:
RuntimeError: If task support is not enabled or task_metadata is missing
Example:
@server.call_tool()
async def handle_tool(name: str, args: dict):
ctx = server.request_context
async def work(task: ServerTaskContext) -> CallToolResult:
result = await task.elicit(
message="Are you sure?",
requestedSchema={"type": "object", ...}
)
confirmed = result.content.get("confirm", False)
return CallToolResult(content=[TextContent(text="Done" if confirmed else "Cancelled")])
return await ctx.experimental.run_task(work)
WARNING: This API is experimental and may change without notice.
"""
if self._task_support is None:
raise RuntimeError("Task support not enabled. Call server.experimental.enable_tasks() first.")
if self._session is None:
raise RuntimeError("Session not available.")
if self.task_metadata is None:
raise RuntimeError(
"Request is not task-augmented (no task field in params). "
"The client must send a task-augmented request."
)
support = self._task_support
# Access task_group via TaskSupport - raises if not in run() context
task_group = support.task_group
task = await support.store.create_task(self.task_metadata, task_id)
task_ctx = ServerTaskContext(
task=task,
store=support.store,
session=self._session,
queue=support.queue,
handler=support.handler,
)
async def execute() -> None:
try:
result = await work(task_ctx)
if not is_terminal(task_ctx.task.status):
await task_ctx.complete(result)
except Exception as e:
if not is_terminal(task_ctx.task.status):
await task_ctx.fail(str(e))
task_group.start_soon(execute)
meta: dict[str, Any] | None = None
if model_immediate_response is not None:
meta = {MODEL_IMMEDIATE_RESPONSE_KEY: model_immediate_response}
return CreateTaskResult(task=task, **{"_meta": meta} if meta else {})
@@ -0,0 +1,220 @@
"""
Experimental server session features for server→client task operations.
This module provides the server-side equivalent of ExperimentalClientFeatures,
allowing the server to send task-augmented requests to the client and poll for results.
WARNING: These APIs are experimental and may change without notice.
"""
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Any, TypeVar
import mcp.types as types
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
from mcp.shared.experimental.tasks.capabilities import (
require_task_augmented_elicitation,
require_task_augmented_sampling,
)
from mcp.shared.experimental.tasks.polling import poll_until_terminal
if TYPE_CHECKING:
from mcp.server.session import ServerSession
ResultT = TypeVar("ResultT", bound=types.Result)
class ExperimentalServerSessionFeatures:
"""
Experimental server session features for server→client task operations.
This provides the server-side equivalent of ExperimentalClientFeatures,
allowing the server to send task-augmented requests to the client and
poll for results.
WARNING: These APIs are experimental and may change without notice.
Access via session.experimental:
result = await session.experimental.elicit_as_task(...)
"""
def __init__(self, session: "ServerSession") -> None:
self._session = session
async def get_task(self, task_id: str) -> types.GetTaskResult:
"""
Send tasks/get to the client to get task status.
Args:
task_id: The task identifier
Returns:
GetTaskResult containing the task status
"""
return await self._session.send_request(
types.ServerRequest(types.GetTaskRequest(params=types.GetTaskRequestParams(taskId=task_id))),
types.GetTaskResult,
)
async def get_task_result(
self,
task_id: str,
result_type: type[ResultT],
) -> ResultT:
"""
Send tasks/result to the client to retrieve the final result.
Args:
task_id: The task identifier
result_type: The expected result type
Returns:
The task result, validated against result_type
"""
return await self._session.send_request(
types.ServerRequest(types.GetTaskPayloadRequest(params=types.GetTaskPayloadRequestParams(taskId=task_id))),
result_type,
)
async def poll_task(self, task_id: str) -> AsyncIterator[types.GetTaskResult]:
"""
Poll a client task until it reaches terminal status.
Yields GetTaskResult for each poll, allowing the caller to react to
status changes. Exits when task reaches a terminal status.
Respects the pollInterval hint from the client.
Args:
task_id: The task identifier
Yields:
GetTaskResult for each poll
"""
async for status in poll_until_terminal(self.get_task, task_id):
yield status
async def elicit_as_task(
self,
message: str,
requestedSchema: types.ElicitRequestedSchema,
*,
ttl: int = 60000,
) -> types.ElicitResult:
"""
Send a task-augmented elicitation to the client and poll until complete.
The client will create a local task, process the elicitation asynchronously,
and return the result when ready. This method handles the full flow:
1. Send elicitation with task field
2. Receive CreateTaskResult from client
3. Poll client's task until terminal
4. Retrieve and return the final ElicitResult
Args:
message: The message to present to the user
requestedSchema: Schema defining the expected response
ttl: Task time-to-live in milliseconds
Returns:
The client's elicitation response
Raises:
McpError: If client doesn't support task-augmented elicitation
"""
client_caps = self._session.client_params.capabilities if self._session.client_params else None
require_task_augmented_elicitation(client_caps)
create_result = await self._session.send_request(
types.ServerRequest(
types.ElicitRequest(
params=types.ElicitRequestFormParams(
message=message,
requestedSchema=requestedSchema,
task=types.TaskMetadata(ttl=ttl),
)
)
),
types.CreateTaskResult,
)
task_id = create_result.task.taskId
async for _ in self.poll_task(task_id):
pass
return await self.get_task_result(task_id, types.ElicitResult)
async def create_message_as_task(
self,
messages: list[types.SamplingMessage],
*,
max_tokens: int,
ttl: int = 60000,
system_prompt: str | None = None,
include_context: types.IncludeContext | None = None,
temperature: float | None = None,
stop_sequences: list[str] | None = None,
metadata: dict[str, Any] | None = None,
model_preferences: types.ModelPreferences | None = None,
tools: list[types.Tool] | None = None,
tool_choice: types.ToolChoice | None = None,
) -> types.CreateMessageResult:
"""
Send a task-augmented sampling request and poll until complete.
The client will create a local task, process the sampling request
asynchronously, and return the result when ready.
Args:
messages: The conversation messages for sampling
max_tokens: Maximum tokens in the response
ttl: Task time-to-live in milliseconds
system_prompt: Optional system prompt
include_context: Context inclusion strategy
temperature: Sampling temperature
stop_sequences: Stop sequences
metadata: Additional metadata
model_preferences: Model selection preferences
tools: Optional list of tools the LLM can use during sampling
tool_choice: Optional control over tool usage behavior
Returns:
The sampling result from the client
Raises:
McpError: If client doesn't support task-augmented sampling or tools
ValueError: If tool_use or tool_result message structure is invalid
"""
client_caps = self._session.client_params.capabilities if self._session.client_params else None
require_task_augmented_sampling(client_caps)
validate_sampling_tools(client_caps, tools, tool_choice)
validate_tool_use_result_messages(messages)
create_result = await self._session.send_request(
types.ServerRequest(
types.CreateMessageRequest(
params=types.CreateMessageRequestParams(
messages=messages,
maxTokens=max_tokens,
systemPrompt=system_prompt,
includeContext=include_context,
temperature=temperature,
stopSequences=stop_sequences,
metadata=metadata,
modelPreferences=model_preferences,
tools=tools,
toolChoice=tool_choice,
task=types.TaskMetadata(ttl=ttl),
)
)
),
types.CreateTaskResult,
)
task_id = create_result.task.taskId
async for _ in self.poll_task(task_id):
pass
return await self.get_task_result(task_id, types.CreateMessageResult)
@@ -0,0 +1,612 @@
"""
ServerTaskContext - Server-integrated task context with elicitation and sampling.
This wraps the pure TaskContext and adds server-specific functionality:
- Elicitation (task.elicit())
- Sampling (task.create_message())
- Status notifications
"""
from typing import Any
import anyio
from mcp.server.experimental.task_result_handler import TaskResultHandler
from mcp.server.session import ServerSession
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.capabilities import (
require_task_augmented_elicitation,
require_task_augmented_sampling,
)
from mcp.shared.experimental.tasks.context import TaskContext
from mcp.shared.experimental.tasks.message_queue import QueuedMessage, TaskMessageQueue
from mcp.shared.experimental.tasks.resolver import Resolver
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.types import (
INVALID_REQUEST,
TASK_STATUS_INPUT_REQUIRED,
TASK_STATUS_WORKING,
ClientCapabilities,
CreateMessageResult,
CreateTaskResult,
ElicitationCapability,
ElicitRequestedSchema,
ElicitResult,
ErrorData,
IncludeContext,
ModelPreferences,
RequestId,
Result,
SamplingCapability,
SamplingMessage,
ServerNotification,
Task,
TaskMetadata,
TaskStatusNotification,
TaskStatusNotificationParams,
Tool,
ToolChoice,
)
class ServerTaskContext:
"""
Server-integrated task context with elicitation and sampling.
This wraps a pure TaskContext and adds server-specific functionality:
- elicit() for sending elicitation requests to the client
- create_message() for sampling requests
- Status notifications via the session
Example:
async def my_task_work(task: ServerTaskContext) -> CallToolResult:
await task.update_status("Starting...")
result = await task.elicit(
message="Continue?",
requestedSchema={"type": "object", "properties": {"ok": {"type": "boolean"}}}
)
if result.content.get("ok"):
return CallToolResult(content=[TextContent(text="Done!")])
else:
return CallToolResult(content=[TextContent(text="Cancelled")])
"""
def __init__(
self,
*,
task: Task,
store: TaskStore,
session: ServerSession,
queue: TaskMessageQueue,
handler: TaskResultHandler | None = None,
):
"""
Create a ServerTaskContext.
Args:
task: The Task object
store: The task store
session: The server session
queue: The message queue for elicitation/sampling
handler: The result handler for response routing (required for elicit/create_message)
"""
self._ctx = TaskContext(task=task, store=store)
self._session = session
self._queue = queue
self._handler = handler
self._store = store
# Delegate pure properties to inner context
@property
def task_id(self) -> str:
"""The task identifier."""
return self._ctx.task_id
@property
def task(self) -> Task:
"""The current task state."""
return self._ctx.task
@property
def is_cancelled(self) -> bool:
"""Whether cancellation has been requested."""
return self._ctx.is_cancelled
def request_cancellation(self) -> None:
"""Request cancellation of this task."""
self._ctx.request_cancellation()
# Enhanced methods with notifications
async def update_status(self, message: str, *, notify: bool = True) -> None:
"""
Update the task's status message.
Args:
message: The new status message
notify: Whether to send a notification to the client
"""
await self._ctx.update_status(message)
if notify:
await self._send_notification()
async def complete(self, result: Result, *, notify: bool = True) -> None:
"""
Mark the task as completed with the given result.
Args:
result: The task result
notify: Whether to send a notification to the client
"""
await self._ctx.complete(result)
if notify:
await self._send_notification()
async def fail(self, error: str, *, notify: bool = True) -> None:
"""
Mark the task as failed with an error message.
Args:
error: The error message
notify: Whether to send a notification to the client
"""
await self._ctx.fail(error)
if notify:
await self._send_notification()
async def _send_notification(self) -> None:
"""Send a task status notification to the client."""
task = self._ctx.task
await self._session.send_notification(
ServerNotification(
TaskStatusNotification(
params=TaskStatusNotificationParams(
taskId=task.taskId,
status=task.status,
statusMessage=task.statusMessage,
createdAt=task.createdAt,
lastUpdatedAt=task.lastUpdatedAt,
ttl=task.ttl,
pollInterval=task.pollInterval,
)
)
)
)
# Server-specific methods: elicitation and sampling
def _check_elicitation_capability(self) -> None:
"""Check if the client supports elicitation."""
if not self._session.check_client_capability(ClientCapabilities(elicitation=ElicitationCapability())):
raise McpError(
ErrorData(
code=INVALID_REQUEST,
message="Client does not support elicitation capability",
)
)
def _check_sampling_capability(self) -> None:
"""Check if the client supports sampling."""
if not self._session.check_client_capability(ClientCapabilities(sampling=SamplingCapability())):
raise McpError(
ErrorData(
code=INVALID_REQUEST,
message="Client does not support sampling capability",
)
)
async def elicit(
self,
message: str,
requestedSchema: ElicitRequestedSchema,
) -> ElicitResult:
"""
Send an elicitation request via the task message queue.
This method:
1. Checks client capability
2. Updates task status to "input_required"
3. Queues the elicitation request
4. Waits for the response (delivered via tasks/result round-trip)
5. Updates task status back to "working"
6. Returns the result
Args:
message: The message to present to the user
requestedSchema: Schema defining the expected response structure
Returns:
The client's response
Raises:
McpError: If client doesn't support elicitation capability
"""
self._check_elicitation_capability()
if self._handler is None:
raise RuntimeError("handler is required for elicit(). Pass handler= to ServerTaskContext.")
# Update status to input_required
await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
# Build the request using session's helper
request = self._session._build_elicit_form_request( # pyright: ignore[reportPrivateUsage]
message=message,
requestedSchema=requestedSchema,
related_task_id=self.task_id,
)
request_id: RequestId = request.id
resolver: Resolver[dict[str, Any]] = Resolver()
self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage]
queued = QueuedMessage(
type="request",
message=request,
resolver=resolver,
original_request_id=request_id,
)
await self._queue.enqueue(self.task_id, queued)
try:
# Wait for response (routed back via TaskResultHandler)
response_data = await resolver.wait()
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
return ElicitResult.model_validate(response_data)
except anyio.get_cancelled_exc_class(): # pragma: no cover
# Coverage can't track async exception handlers reliably.
# This path is tested in test_elicit_restores_status_on_cancellation
# which verifies status is restored to "working" after cancellation.
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
raise
async def elicit_url(
self,
message: str,
url: str,
elicitation_id: str,
) -> ElicitResult:
"""
Send a URL mode elicitation request via the task message queue.
This directs the user to an external URL for out-of-band interactions
like OAuth flows, credential collection, or payment processing.
This method:
1. Checks client capability
2. Updates task status to "input_required"
3. Queues the elicitation request
4. Waits for the response (delivered via tasks/result round-trip)
5. Updates task status back to "working"
6. Returns the result
Args:
message: Human-readable explanation of why the interaction is needed
url: The URL the user should navigate to
elicitation_id: Unique identifier for tracking this elicitation
Returns:
The client's response indicating acceptance, decline, or cancellation
Raises:
McpError: If client doesn't support elicitation capability
RuntimeError: If handler is not configured
"""
self._check_elicitation_capability()
if self._handler is None:
raise RuntimeError("handler is required for elicit_url(). Pass handler= to ServerTaskContext.")
# Update status to input_required
await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
# Build the request using session's helper
request = self._session._build_elicit_url_request( # pyright: ignore[reportPrivateUsage]
message=message,
url=url,
elicitation_id=elicitation_id,
related_task_id=self.task_id,
)
request_id: RequestId = request.id
resolver: Resolver[dict[str, Any]] = Resolver()
self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage]
queued = QueuedMessage(
type="request",
message=request,
resolver=resolver,
original_request_id=request_id,
)
await self._queue.enqueue(self.task_id, queued)
try:
# Wait for response (routed back via TaskResultHandler)
response_data = await resolver.wait()
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
return ElicitResult.model_validate(response_data)
except anyio.get_cancelled_exc_class(): # pragma: no cover
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
raise
async def create_message(
self,
messages: list[SamplingMessage],
*,
max_tokens: int,
system_prompt: str | None = None,
include_context: IncludeContext | None = None,
temperature: float | None = None,
stop_sequences: list[str] | None = None,
metadata: dict[str, Any] | None = None,
model_preferences: ModelPreferences | None = None,
tools: list[Tool] | None = None,
tool_choice: ToolChoice | None = None,
) -> CreateMessageResult:
"""
Send a sampling request via the task message queue.
This method:
1. Checks client capability
2. Updates task status to "input_required"
3. Queues the sampling request
4. Waits for the response (delivered via tasks/result round-trip)
5. Updates task status back to "working"
6. Returns the result
Args:
messages: The conversation messages for sampling
max_tokens: Maximum tokens in the response
system_prompt: Optional system prompt
include_context: Context inclusion strategy
temperature: Sampling temperature
stop_sequences: Stop sequences
metadata: Additional metadata
model_preferences: Model selection preferences
tools: Optional list of tools the LLM can use during sampling
tool_choice: Optional control over tool usage behavior
Returns:
The sampling result from the client
Raises:
McpError: If client doesn't support sampling capability or tools
ValueError: If tool_use or tool_result message structure is invalid
"""
self._check_sampling_capability()
client_caps = self._session.client_params.capabilities if self._session.client_params else None
validate_sampling_tools(client_caps, tools, tool_choice)
validate_tool_use_result_messages(messages)
if self._handler is None:
raise RuntimeError("handler is required for create_message(). Pass handler= to ServerTaskContext.")
# Update status to input_required
await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
# Build the request using session's helper
request = self._session._build_create_message_request( # pyright: ignore[reportPrivateUsage]
messages=messages,
max_tokens=max_tokens,
system_prompt=system_prompt,
include_context=include_context,
temperature=temperature,
stop_sequences=stop_sequences,
metadata=metadata,
model_preferences=model_preferences,
tools=tools,
tool_choice=tool_choice,
related_task_id=self.task_id,
)
request_id: RequestId = request.id
resolver: Resolver[dict[str, Any]] = Resolver()
self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage]
queued = QueuedMessage(
type="request",
message=request,
resolver=resolver,
original_request_id=request_id,
)
await self._queue.enqueue(self.task_id, queued)
try:
# Wait for response (routed back via TaskResultHandler)
response_data = await resolver.wait()
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
return CreateMessageResult.model_validate(response_data)
except anyio.get_cancelled_exc_class(): # pragma: no cover
# Coverage can't track async exception handlers reliably.
# This path is tested in test_create_message_restores_status_on_cancellation
# which verifies status is restored to "working" after cancellation.
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
raise
async def elicit_as_task(
self,
message: str,
requestedSchema: ElicitRequestedSchema,
*,
ttl: int = 60000,
) -> ElicitResult:
"""
Send a task-augmented elicitation via the queue, then poll client.
This is for use inside a task-augmented tool call when you want the client
to handle the elicitation as its own task. The elicitation request is queued
and delivered when the client calls tasks/result. After the client responds
with CreateTaskResult, we poll the client's task until complete.
Args:
message: The message to present to the user
requestedSchema: Schema defining the expected response structure
ttl: Task time-to-live in milliseconds for the client's task
Returns:
The client's elicitation response
Raises:
McpError: If client doesn't support task-augmented elicitation
RuntimeError: If handler is not configured
"""
client_caps = self._session.client_params.capabilities if self._session.client_params else None
require_task_augmented_elicitation(client_caps)
if self._handler is None:
raise RuntimeError("handler is required for elicit_as_task()")
# Update status to input_required
await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
request = self._session._build_elicit_form_request( # pyright: ignore[reportPrivateUsage]
message=message,
requestedSchema=requestedSchema,
related_task_id=self.task_id,
task=TaskMetadata(ttl=ttl),
)
request_id: RequestId = request.id
resolver: Resolver[dict[str, Any]] = Resolver()
self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage]
queued = QueuedMessage(
type="request",
message=request,
resolver=resolver,
original_request_id=request_id,
)
await self._queue.enqueue(self.task_id, queued)
try:
# Wait for initial response (CreateTaskResult from client)
response_data = await resolver.wait()
create_result = CreateTaskResult.model_validate(response_data)
client_task_id = create_result.task.taskId
# Poll the client's task using session.experimental
async for _ in self._session.experimental.poll_task(client_task_id):
pass
# Get final result from client
result = await self._session.experimental.get_task_result(
client_task_id,
ElicitResult,
)
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
return result
except anyio.get_cancelled_exc_class(): # pragma: no cover
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
raise
async def create_message_as_task(
self,
messages: list[SamplingMessage],
*,
max_tokens: int,
ttl: int = 60000,
system_prompt: str | None = None,
include_context: IncludeContext | None = None,
temperature: float | None = None,
stop_sequences: list[str] | None = None,
metadata: dict[str, Any] | None = None,
model_preferences: ModelPreferences | None = None,
tools: list[Tool] | None = None,
tool_choice: ToolChoice | None = None,
) -> CreateMessageResult:
"""
Send a task-augmented sampling request via the queue, then poll client.
This is for use inside a task-augmented tool call when you want the client
to handle the sampling as its own task. The request is queued and delivered
when the client calls tasks/result. After the client responds with
CreateTaskResult, we poll the client's task until complete.
Args:
messages: The conversation messages for sampling
max_tokens: Maximum tokens in the response
ttl: Task time-to-live in milliseconds for the client's task
system_prompt: Optional system prompt
include_context: Context inclusion strategy
temperature: Sampling temperature
stop_sequences: Stop sequences
metadata: Additional metadata
model_preferences: Model selection preferences
tools: Optional list of tools the LLM can use during sampling
tool_choice: Optional control over tool usage behavior
Returns:
The sampling result from the client
Raises:
McpError: If client doesn't support task-augmented sampling or tools
ValueError: If tool_use or tool_result message structure is invalid
RuntimeError: If handler is not configured
"""
client_caps = self._session.client_params.capabilities if self._session.client_params else None
require_task_augmented_sampling(client_caps)
validate_sampling_tools(client_caps, tools, tool_choice)
validate_tool_use_result_messages(messages)
if self._handler is None:
raise RuntimeError("handler is required for create_message_as_task()")
# Update status to input_required
await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
# Build request WITH task field for task-augmented sampling
request = self._session._build_create_message_request( # pyright: ignore[reportPrivateUsage]
messages=messages,
max_tokens=max_tokens,
system_prompt=system_prompt,
include_context=include_context,
temperature=temperature,
stop_sequences=stop_sequences,
metadata=metadata,
model_preferences=model_preferences,
tools=tools,
tool_choice=tool_choice,
related_task_id=self.task_id,
task=TaskMetadata(ttl=ttl),
)
request_id: RequestId = request.id
resolver: Resolver[dict[str, Any]] = Resolver()
self._handler._pending_requests[request_id] = resolver # pyright: ignore[reportPrivateUsage]
queued = QueuedMessage(
type="request",
message=request,
resolver=resolver,
original_request_id=request_id,
)
await self._queue.enqueue(self.task_id, queued)
try:
# Wait for initial response (CreateTaskResult from client)
response_data = await resolver.wait()
create_result = CreateTaskResult.model_validate(response_data)
client_task_id = create_result.task.taskId
# Poll the client's task using session.experimental
async for _ in self._session.experimental.poll_task(client_task_id):
pass
# Get final result from client
result = await self._session.experimental.get_task_result(
client_task_id,
CreateMessageResult,
)
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
return result
except anyio.get_cancelled_exc_class(): # pragma: no cover
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
raise
@@ -0,0 +1,235 @@
"""
TaskResultHandler - Integrated handler for tasks/result endpoint.
This implements the dequeue-send-wait pattern from the MCP Tasks spec:
1. Dequeue all pending messages for the task
2. Send them to the client via transport with relatedRequestId routing
3. Wait if task is not in terminal state
4. Return final result when task completes
This is the core of the task message queue pattern.
"""
import logging
from typing import Any
import anyio
from mcp.server.session import ServerSession
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY, is_terminal
from mcp.shared.experimental.tasks.message_queue import TaskMessageQueue
from mcp.shared.experimental.tasks.resolver import Resolver
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.types import (
INVALID_PARAMS,
ErrorData,
GetTaskPayloadRequest,
GetTaskPayloadResult,
JSONRPCMessage,
RelatedTaskMetadata,
RequestId,
)
logger = logging.getLogger(__name__)
class TaskResultHandler:
"""
Handler for tasks/result that implements the message queue pattern.
This handler:
1. Dequeues pending messages (elicitations, notifications) for the task
2. Sends them to the client via the response stream
3. Waits for responses and resolves them back to callers
4. Blocks until task reaches terminal state
5. Returns the final result
Usage:
# Create handler with store and queue
handler = TaskResultHandler(task_store, message_queue)
# Register it with the server
@server.experimental.get_task_result()
async def handle_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
ctx = server.request_context
return await handler.handle(req, ctx.session, ctx.request_id)
# Or use the convenience method
handler.register(server)
"""
def __init__(
self,
store: TaskStore,
queue: TaskMessageQueue,
):
self._store = store
self._queue = queue
# Map from internal request ID to resolver for routing responses
self._pending_requests: dict[RequestId, Resolver[dict[str, Any]]] = {}
async def send_message(
self,
session: ServerSession,
message: SessionMessage,
) -> None:
"""
Send a message via the session.
This is a helper for delivering queued task messages.
"""
await session.send_message(message)
async def handle(
self,
request: GetTaskPayloadRequest,
session: ServerSession,
request_id: RequestId,
) -> GetTaskPayloadResult:
"""
Handle a tasks/result request.
This implements the dequeue-send-wait loop:
1. Dequeue all pending messages
2. Send each via transport with relatedRequestId = this request's ID
3. If task not terminal, wait for status change
4. Loop until task is terminal
5. Return final result
Args:
request: The GetTaskPayloadRequest
session: The server session for sending messages
request_id: The request ID for relatedRequestId routing
Returns:
GetTaskPayloadResult with the task's final payload
"""
task_id = request.params.taskId
while True:
task = await self._store.get_task(task_id)
if task is None:
raise McpError(
ErrorData(
code=INVALID_PARAMS,
message=f"Task not found: {task_id}",
)
)
await self._deliver_queued_messages(task_id, session, request_id)
# If task is terminal, return result
if is_terminal(task.status):
result = await self._store.get_result(task_id)
# GetTaskPayloadResult is a Result with extra="allow"
# The stored result contains the actual payload data
# Per spec: tasks/result MUST include _meta with related-task metadata
related_task = RelatedTaskMetadata(taskId=task_id)
related_task_meta: dict[str, Any] = {RELATED_TASK_METADATA_KEY: related_task.model_dump(by_alias=True)}
if result is not None:
result_data = result.model_dump(by_alias=True)
existing_meta: dict[str, Any] = result_data.get("_meta") or {}
result_data["_meta"] = {**existing_meta, **related_task_meta}
return GetTaskPayloadResult.model_validate(result_data)
return GetTaskPayloadResult.model_validate({"_meta": related_task_meta})
# Wait for task update (status change or new messages)
await self._wait_for_task_update(task_id)
async def _deliver_queued_messages(
self,
task_id: str,
session: ServerSession,
request_id: RequestId,
) -> None:
"""
Dequeue and send all pending messages for a task.
Each message is sent via the session's write stream with
relatedRequestId set so responses route back to this stream.
"""
while True:
message = await self._queue.dequeue(task_id)
if message is None:
break
# If this is a request (not notification), wait for response
if message.type == "request" and message.resolver is not None:
# Store the resolver so we can route the response back
original_id = message.original_request_id
if original_id is not None:
self._pending_requests[original_id] = message.resolver
logger.debug("Delivering queued message for task %s: %s", task_id, message.type)
# Send the message with relatedRequestId for routing
session_message = SessionMessage(
message=JSONRPCMessage(message.message),
metadata=ServerMessageMetadata(related_request_id=request_id),
)
await self.send_message(session, session_message)
async def _wait_for_task_update(self, task_id: str) -> None:
"""
Wait for task to be updated (status change or new message).
Races between store update and queue message - first one wins.
"""
async with anyio.create_task_group() as tg:
async def wait_for_store() -> None:
try:
await self._store.wait_for_update(task_id)
except Exception:
pass
finally:
tg.cancel_scope.cancel()
async def wait_for_queue() -> None:
try:
await self._queue.wait_for_message(task_id)
except Exception:
pass
finally:
tg.cancel_scope.cancel()
tg.start_soon(wait_for_store)
tg.start_soon(wait_for_queue)
def route_response(self, request_id: RequestId, response: dict[str, Any]) -> bool:
"""
Route a response back to the waiting resolver.
This is called when a response arrives for a queued request.
Args:
request_id: The request ID from the response
response: The response data
Returns:
True if response was routed, False if no pending request
"""
resolver = self._pending_requests.pop(request_id, None)
if resolver is not None and not resolver.done():
resolver.set_result(response)
return True
return False
def route_error(self, request_id: RequestId, error: ErrorData) -> bool:
"""
Route an error back to the waiting resolver.
Args:
request_id: The request ID from the error response
error: The error data
Returns:
True if error was routed, False if no pending request
"""
resolver = self._pending_requests.pop(request_id, None)
if resolver is not None and not resolver.done():
resolver.set_exception(McpError(error))
return True
return False
@@ -0,0 +1,115 @@
"""
TaskSupport - Configuration for experimental task support.
This module provides the TaskSupport class which encapsulates all the
infrastructure needed for task-augmented requests: store, queue, and handler.
"""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
import anyio
from anyio.abc import TaskGroup
from mcp.server.experimental.task_result_handler import TaskResultHandler
from mcp.server.session import ServerSession
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue
from mcp.shared.experimental.tasks.store import TaskStore
@dataclass
class TaskSupport:
"""
Configuration for experimental task support.
Encapsulates the task store, message queue, result handler, and task group
for spawning background work.
When enabled on a server, this automatically:
- Configures response routing for each session
- Provides default handlers for task operations
- Manages a task group for background task execution
Example:
# Simple in-memory setup
server.experimental.enable_tasks()
# Custom store/queue for distributed systems
server.experimental.enable_tasks(
store=RedisTaskStore(redis_url),
queue=RedisTaskMessageQueue(redis_url),
)
"""
store: TaskStore
queue: TaskMessageQueue
handler: TaskResultHandler = field(init=False)
_task_group: TaskGroup | None = field(init=False, default=None)
def __post_init__(self) -> None:
"""Create the result handler from store and queue."""
self.handler = TaskResultHandler(self.store, self.queue)
@property
def task_group(self) -> TaskGroup:
"""Get the task group for spawning background work.
Raises:
RuntimeError: If not within a run() context
"""
if self._task_group is None:
raise RuntimeError("TaskSupport not running. Ensure Server.run() is active.")
return self._task_group
@asynccontextmanager
async def run(self) -> AsyncIterator[None]:
"""
Run the task support lifecycle.
This creates a task group for spawning background task work.
Called automatically by Server.run().
Usage:
async with task_support.run():
# Task group is now available
...
"""
async with anyio.create_task_group() as tg:
self._task_group = tg
try:
yield
finally:
self._task_group = None
def configure_session(self, session: ServerSession) -> None:
"""
Configure a session for task support.
This registers the result handler as a response router so that
responses to queued requests (elicitation, sampling) are routed
back to the waiting resolvers.
Called automatically by Server.run() for each new session.
Args:
session: The session to configure
"""
session.add_response_router(self.handler)
@classmethod
def in_memory(cls) -> "TaskSupport":
"""
Create in-memory task support.
Suitable for development, testing, and single-process servers.
For distributed systems, provide custom store and queue implementations.
Returns:
TaskSupport configured with in-memory store and queue
"""
return cls(
store=InMemoryTaskStore(),
queue=InMemoryTaskMessageQueue(),
)
@@ -0,0 +1,11 @@
"""FastMCP - A more ergonomic interface for MCP servers."""
from importlib.metadata import version
from mcp.types import Icon
from .server import Context, FastMCP
from .utilities.types import Audio, Image
__version__ = version("mcp")
__all__ = ["FastMCP", "Context", "Image", "Audio", "Icon"]
@@ -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."""
@@ -0,0 +1,4 @@
from .base import Prompt
from .manager import PromptManager
__all__ = ["Prompt", "PromptManager"]
@@ -0,0 +1,183 @@
"""Base classes for FastMCP prompts."""
from __future__ import annotations
import inspect
from collections.abc import Awaitable, Callable, Sequence
from typing import TYPE_CHECKING, Any, Literal
import pydantic_core
from pydantic import BaseModel, Field, TypeAdapter, validate_call
from mcp.server.fastmcp.utilities.context_injection import find_context_parameter, inject_context
from mcp.server.fastmcp.utilities.func_metadata import func_metadata
from mcp.types import ContentBlock, Icon, TextContent
if TYPE_CHECKING:
from mcp.server.fastmcp.server import Context
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT, RequestT
class Message(BaseModel):
"""Base class for all prompt messages."""
role: Literal["user", "assistant"]
content: ContentBlock
def __init__(self, content: str | ContentBlock, **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 | ContentBlock, **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 | ContentBlock, **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")
title: str | None = Field(None, description="Human-readable title 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)
icons: list[Icon] | None = Field(default=None, description="Optional list of icons for this prompt")
context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context", exclude=True)
@classmethod
def from_function(
cls,
fn: Callable[..., PromptResult | Awaitable[PromptResult]],
name: str | None = None,
title: str | None = None,
description: str | None = None,
icons: list[Icon] | None = None,
context_kwarg: 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 == "<lambda>": # pragma: no cover
raise ValueError("You must provide a name for lambda functions")
# Find context parameter if it exists
if context_kwarg is None: # pragma: no branch
context_kwarg = find_context_parameter(fn)
# Get schema from func_metadata, excluding context parameter
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()
# Convert parameters to PromptArguments
arguments: list[PromptArgument] = []
if "properties" in parameters: # pragma: no branch
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,
title=title,
description=description or fn.__doc__ or "",
arguments=arguments,
fn=fn,
icons=icons,
context_kwarg=context_kwarg,
)
async def render(
self,
arguments: dict[str, Any] | None = None,
context: Context[ServerSessionT, LifespanContextT, RequestT] | 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:
# Add context to arguments if needed
call_args = inject_context(self.fn, arguments or {}, context, self.context_kwarg)
# Call function and check if result is a coroutine
result = self.fn(**call_args)
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: # pragma: no cover
content = pydantic_core.to_json(msg, fallback=str, indent=2).decode()
messages.append(Message(role="user", content=content))
except Exception: # pragma: no cover
raise ValueError(f"Could not convert prompt result to message: {msg}")
return messages
except Exception as e: # pragma: no cover
raise ValueError(f"Error rendering prompt {self.name}: {e}")
@@ -0,0 +1,60 @@
"""Prompt management functionality."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from mcp.server.fastmcp.prompts.base import Message, Prompt
from mcp.server.fastmcp.utilities.logging import get_logger
if TYPE_CHECKING:
from mcp.server.fastmcp.server import Context
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT, RequestT
logger = get_logger(__name__)
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,
context: Context[ServerSessionT, LifespanContextT, RequestT] | 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, context=context)
@@ -0,0 +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__ = [
"Resource",
"TextResource",
"BinaryResource",
"FunctionResource",
"FileResource",
"HttpResource",
"DirectoryResource",
"ResourceTemplate",
"ResourceManager",
]
@@ -0,0 +1,50 @@
"""Base classes and interfaces for FastMCP resources."""
import abc
from typing import Annotated, Any
from pydantic import (
AnyUrl,
BaseModel,
ConfigDict,
Field,
UrlConstraints,
ValidationInfo,
field_validator,
)
from mcp.types import Annotations, Icon
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)
title: str | None = Field(description="Human-readable title 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\-+.]+(;\s*[a-zA-Z0-9\-_.]+=[a-zA-Z0-9\-_.]+)*$",
)
icons: list[Icon] | None = Field(default=None, description="Optional list of icons for this resource")
annotations: Annotations | None = Field(default=None, description="Optional annotations for the resource")
meta: dict[str, Any] | None = Field(default=None, description="Optional metadata for this resource")
@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 # pragma: no cover
@@ -0,0 +1,115 @@
"""Resource manager functionality."""
from __future__ import annotations
from collections.abc import Callable
from typing import TYPE_CHECKING, Any
from pydantic import AnyUrl
from mcp.server.fastmcp.resources.base import Resource
from mcp.server.fastmcp.resources.templates import ResourceTemplate
from mcp.server.fastmcp.utilities.logging import get_logger
from mcp.types import Annotations, Icon
if TYPE_CHECKING:
from mcp.server.fastmcp.server import Context
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT, RequestT
logger = get_logger(__name__)
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,
title: str | None = None,
description: str | None = None,
mime_type: str | None = None,
icons: list[Icon] | None = None,
annotations: Annotations | None = None,
meta: dict[str, Any] | None = None,
) -> ResourceTemplate:
"""Add a template from a function."""
template = ResourceTemplate.from_function(
fn,
uri_template=uri_template,
name=name,
title=title,
description=description,
mime_type=mime_type,
icons=icons,
annotations=annotations,
meta=meta,
)
self._templates[template.uri_template] = template
return template
async def get_resource(
self,
uri: AnyUrl | str,
context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
) -> 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, context=context)
except Exception as e: # pragma: no cover
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())
@@ -0,0 +1,122 @@
"""Resource template functionality."""
from __future__ import annotations
import inspect
import re
from collections.abc import Callable
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field, validate_call
from mcp.server.fastmcp.resources.types import FunctionResource, Resource
from mcp.server.fastmcp.utilities.context_injection import find_context_parameter, inject_context
from mcp.server.fastmcp.utilities.func_metadata import func_metadata
from mcp.types import Annotations, Icon
if TYPE_CHECKING:
from mcp.server.fastmcp.server import Context
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT, RequestT
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")
title: str | None = Field(description="Human-readable title of the resource", default=None)
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")
icons: list[Icon] | None = Field(default=None, description="Optional list of icons for the resource template")
annotations: Annotations | None = Field(default=None, description="Optional annotations for the resource template")
meta: dict[str, Any] | None = Field(default=None, description="Optional metadata for this resource template")
fn: Callable[..., Any] = Field(exclude=True)
parameters: dict[str, Any] = Field(description="JSON schema for function parameters")
context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context")
@classmethod
def from_function(
cls,
fn: Callable[..., Any],
uri_template: str,
name: str | None = None,
title: str | None = None,
description: str | None = None,
mime_type: str | None = None,
icons: list[Icon] | None = None,
annotations: Annotations | None = None,
meta: dict[str, Any] | None = None,
context_kwarg: str | None = None,
) -> ResourceTemplate:
"""Create a template from a function."""
func_name = name or fn.__name__
if func_name == "<lambda>":
raise ValueError("You must provide a name for lambda functions") # pragma: no cover
# Find context parameter if it exists
if context_kwarg is None: # pragma: no branch
context_kwarg = find_context_parameter(fn)
# Get schema from func_metadata, excluding context parameter
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()
# ensure the arguments are properly cast
fn = validate_call(fn)
return cls(
uri_template=uri_template,
name=func_name,
title=title,
description=description or fn.__doc__ or "",
mime_type=mime_type or "text/plain",
icons=icons,
annotations=annotations,
meta=meta,
fn=fn,
parameters=parameters,
context_kwarg=context_kwarg,
)
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],
context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
) -> Resource:
"""Create a resource from the template with the given parameters."""
try:
# Add context to params if needed
params = inject_context(self.fn, params, context, self.context_kwarg)
# 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,
title=self.title,
description=self.description,
mime_type=self.mime_type,
icons=self.icons,
annotations=self.annotations,
meta=self.meta,
fn=lambda: result, # Capture result in closure
)
except Exception as e:
raise ValueError(f"Error creating resource from template: {e}")
@@ -0,0 +1,203 @@
"""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
import pydantic_core
from pydantic import AnyUrl, Field, ValidationInfo, validate_call
from mcp.server.fastmcp.resources.base import Resource
from mcp.types import Annotations, Icon
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 # pragma: no cover
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 # pragma: no cover
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:
# Call the function first to see if it returns a coroutine
result = self.fn()
# If it's a coroutine, await it
if inspect.iscoroutine(result):
result = await result
if isinstance(result, Resource): # pragma: no cover
return await result.read()
elif isinstance(result, bytes):
return result
elif isinstance(result, str):
return result
else:
return pydantic_core.to_json(result, fallback=str, indent=2).decode()
except Exception as e:
raise ValueError(f"Error reading resource {self.uri}: {e}")
@classmethod
def from_function(
cls,
fn: Callable[..., Any],
uri: str,
name: str | None = None,
title: str | None = None,
description: str | None = None,
mime_type: str | None = None,
icons: list[Icon] | None = None,
annotations: Annotations | None = None,
meta: dict[str, Any] | None = None,
) -> "FunctionResource":
"""Create a FunctionResource from a function."""
func_name = name or fn.__name__
if func_name == "<lambda>": # pragma: no cover
raise ValueError("You must provide a name for lambda functions")
# ensure the arguments are properly cast
fn = validate_call(fn)
return cls(
uri=AnyUrl(uri),
name=func_name,
title=title,
description=description or fn.__doc__ or "",
mime_type=mime_type or "text/plain",
fn=fn,
icons=icons,
annotations=annotations,
meta=meta,
)
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: # pragma: no cover
"""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: # pragma: no cover
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: # pragma: no cover
"""Ensure path is absolute."""
if not path.is_absolute():
raise ValueError("Path must be absolute")
return path
def list_files(self) -> list[Path]: # pragma: no cover
"""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 # pragma: no cover
"""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}")
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,4 @@
from .base import Tool
from .tool_manager import ToolManager
__all__ = ["Tool", "ToolManager"]
@@ -0,0 +1,126 @@
from __future__ import annotations as _annotations
import functools
import inspect
from collections.abc import Callable
from functools import cached_property
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from mcp.server.fastmcp.exceptions import ToolError
from mcp.server.fastmcp.utilities.context_injection import find_context_parameter
from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata, func_metadata
from mcp.shared.exceptions import UrlElicitationRequiredError
from mcp.shared.tool_name_validation import validate_and_warn_tool_name
from mcp.types import Icon, ToolAnnotations
if TYPE_CHECKING:
from mcp.server.fastmcp.server import Context
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT, RequestT
class Tool(BaseModel):
"""Internal tool registration info."""
fn: Callable[..., Any] = Field(exclude=True)
name: str = Field(description="Name of the tool")
title: str | None = Field(None, description="Human-readable title 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")
annotations: ToolAnnotations | None = Field(None, description="Optional annotations for the tool")
icons: list[Icon] | None = Field(default=None, description="Optional list of icons for this tool")
meta: dict[str, Any] | None = Field(default=None, description="Optional metadata for this tool")
@cached_property
def output_schema(self) -> dict[str, Any] | None:
return self.fn_metadata.output_schema
@classmethod
def from_function(
cls,
fn: Callable[..., Any],
name: str | None = None,
title: str | None = None,
description: str | None = None,
context_kwarg: str | None = None,
annotations: ToolAnnotations | None = None,
icons: list[Icon] | None = None,
meta: dict[str, Any] | None = None,
structured_output: bool | None = None,
) -> Tool:
"""Create a Tool from a function."""
func_name = name or fn.__name__
validate_and_warn_tool_name(func_name)
if func_name == "<lambda>":
raise ValueError("You must provide a name for lambda functions")
func_doc = description or fn.__doc__ or ""
is_async = _is_async_callable(fn)
if context_kwarg is None: # pragma: no branch
context_kwarg = find_context_parameter(fn)
func_arg_metadata = func_metadata(
fn,
skip_names=[context_kwarg] if context_kwarg is not None else [],
structured_output=structured_output,
)
parameters = func_arg_metadata.arg_model.model_json_schema(by_alias=True)
return cls(
fn=fn,
name=func_name,
title=title,
description=func_doc,
parameters=parameters,
fn_metadata=func_arg_metadata,
is_async=is_async,
context_kwarg=context_kwarg,
annotations=annotations,
icons=icons,
meta=meta,
)
async def run(
self,
arguments: dict[str, Any],
context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
convert_result: bool = False,
) -> Any:
"""Run the tool with arguments."""
try:
result = 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,
)
if convert_result:
result = self.fn_metadata.convert_result(result)
return result
except UrlElicitationRequiredError:
# Re-raise UrlElicitationRequiredError so it can be properly handled
# as an MCP error response with code -32042
raise
except Exception as e:
raise ToolError(f"Error executing tool {self.name}: {e}") from e
def _is_async_callable(obj: Any) -> bool:
while isinstance(obj, functools.partial): # pragma: no cover
obj = obj.func
return inspect.iscoroutinefunction(obj) or (
callable(obj) and inspect.iscoroutinefunction(getattr(obj, "__call__", None))
)
@@ -0,0 +1,93 @@
from __future__ import annotations as _annotations
from collections.abc import Callable
from typing import TYPE_CHECKING, Any
from mcp.server.fastmcp.exceptions import ToolError
from mcp.server.fastmcp.tools.base import Tool
from mcp.server.fastmcp.utilities.logging import get_logger
from mcp.shared.context import LifespanContextT, RequestT
from mcp.types import Icon, ToolAnnotations
if TYPE_CHECKING:
from mcp.server.fastmcp.server import Context
from mcp.server.session import ServerSessionT
logger = get_logger(__name__)
class ToolManager:
"""Manages FastMCP tools."""
def __init__(
self,
warn_on_duplicate_tools: bool = True,
*,
tools: list[Tool] | None = None,
):
self._tools: dict[str, Tool] = {}
if tools is not None:
for tool in tools:
if warn_on_duplicate_tools and tool.name in self._tools:
logger.warning(f"Tool already exists: {tool.name}")
self._tools[tool.name] = 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,
title: str | None = None,
description: str | None = None,
annotations: ToolAnnotations | None = None,
icons: list[Icon] | None = None,
meta: dict[str, Any] | None = None,
structured_output: bool | None = None,
) -> Tool:
"""Add a tool to the server."""
tool = Tool.from_function(
fn,
name=name,
title=title,
description=description,
annotations=annotations,
icons=icons,
meta=meta,
structured_output=structured_output,
)
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
def remove_tool(self, name: str) -> None:
"""Remove a tool by name."""
if name not in self._tools:
raise ToolError(f"Unknown tool: {name}")
del self._tools[name]
async def call_tool(
self,
name: str,
arguments: dict[str, Any],
context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
convert_result: bool = False,
) -> 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, convert_result=convert_result)
@@ -0,0 +1 @@
"""FastMCP utility modules."""
@@ -0,0 +1,68 @@
"""Context injection utilities for FastMCP."""
from __future__ import annotations
import inspect
import typing
from collections.abc import Callable
from typing import Any
def find_context_parameter(fn: Callable[..., Any]) -> str | None:
"""Find the parameter that should receive the Context object.
Searches through the function's signature to find a parameter
with a Context type annotation.
Args:
fn: The function to inspect
Returns:
The name of the context parameter, or None if not found
"""
from mcp.server.fastmcp.server import Context
# Get type hints to properly resolve string annotations
try:
hints = typing.get_type_hints(fn)
except Exception:
# If we can't resolve type hints, we can't find the context parameter
return None
# Check each parameter's type hint
for param_name, annotation in hints.items():
# Handle direct Context type
if inspect.isclass(annotation) and issubclass(annotation, Context):
return param_name
# Handle generic types like Optional[Context]
origin = typing.get_origin(annotation)
if origin is not None:
args = typing.get_args(annotation)
for arg in args:
if inspect.isclass(arg) and issubclass(arg, Context):
return param_name
return None
def inject_context(
fn: Callable[..., Any],
kwargs: dict[str, Any],
context: Any | None,
context_kwarg: str | None,
) -> dict[str, Any]:
"""Inject context into function kwargs if needed.
Args:
fn: The function that will be called
kwargs: The current keyword arguments
context: The context object to inject (if any)
context_kwarg: The name of the parameter to inject into
Returns:
Updated kwargs with context injected if applicable
"""
if context_kwarg is not None and context is not None:
return {**kwargs, context_kwarg: context}
return kwargs
@@ -0,0 +1,533 @@
import inspect
import json
from collections.abc import Awaitable, Callable, Sequence
from itertools import chain
from types import GenericAlias
from typing import Annotated, Any, cast, get_args, get_origin, get_type_hints
import pydantic_core
from pydantic import (
BaseModel,
ConfigDict,
Field,
RootModel,
WithJsonSchema,
create_model,
)
from pydantic.fields import FieldInfo
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind
from typing_extensions import is_typeddict
from typing_inspection.introspection import (
UNKNOWN,
AnnotationSource,
ForbiddenQualifier,
inspect_annotation,
is_union_origin,
)
from mcp.server.fastmcp.exceptions import InvalidSignature
from mcp.server.fastmcp.utilities.logging import get_logger
from mcp.server.fastmcp.utilities.types import Audio, Image
from mcp.types import CallToolResult, ContentBlock, TextContent
logger = get_logger(__name__)
class StrictJsonSchema(GenerateJsonSchema):
"""A JSON schema generator that raises exceptions instead of emitting warnings.
This is used to detect non-serializable types during schema generation.
"""
def emit_warning(self, kind: JsonSchemaWarningKind, detail: str) -> None:
# Raise an exception instead of emitting a warning
raise ValueError(f"JSON schema warning: {kind} - {detail}")
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, field_info in self.__class__.model_fields.items():
value = getattr(self, field_name)
# Use the alias if it exists, otherwise use the field name
output_name = field_info.alias if field_info.alias else field_name
kwargs[output_name] = value
return kwargs
model_config = ConfigDict(
arbitrary_types_allowed=True,
)
class FuncMetadata(BaseModel):
arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
output_schema: dict[str, Any] | None = None
output_model: Annotated[type[BaseModel], WithJsonSchema(None)] | None = None
wrap_output: bool = False
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:
return await fn(**arguments_parsed_dict)
else:
return fn(**arguments_parsed_dict)
def convert_result(self, result: Any) -> Any:
"""
Convert the result of a function call to the appropriate format for
the lowlevel server tool call handler:
- If output_model is None, return the unstructured content directly.
- If output_model is not None, convert the result to structured output format
(dict[str, Any]) and return both unstructured and structured content.
Note: we return unstructured content here **even though the lowlevel server
tool call handler provides generic backwards compatibility serialization of
structured content**. This is for FastMCP backwards compatibility: we need to
retain FastMCP's ad hoc conversion logic for constructing unstructured output
from function return values, whereas the lowlevel server simply serializes
the structured output.
"""
if isinstance(result, CallToolResult):
if self.output_schema is not None:
assert self.output_model is not None, "Output model must be set if output schema is defined"
self.output_model.model_validate(result.structuredContent)
return result
unstructured_content = _convert_to_content(result)
if self.output_schema is None:
return unstructured_content
else:
if self.wrap_output:
result = {"result": result}
assert self.output_model is not None, "Output model must be set if output schema is defined"
validated = self.output_model.model_validate(result)
structured_content = validated.model_dump(mode="json", by_alias=True)
return (unstructured_content, structured_content)
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
# Build a mapping from input keys (including aliases) to field info
key_to_field_info: dict[str, FieldInfo] = {}
for field_name, field_info in self.arg_model.model_fields.items():
# Map both the field name and its alias (if any) to the field info
key_to_field_info[field_name] = field_info
if field_info.alias:
key_to_field_info[field_info.alias] = field_info
for data_key, data_value in data.items():
if data_key not in key_to_field_info: # pragma: no cover
continue
field_info = key_to_field_info[data_key]
if isinstance(data_value, str) and field_info.annotation is not str:
try:
pre_parsed = json.loads(data_value)
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[data_key] = 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] = (),
structured_output: bool | None = None,
) -> FuncMetadata:
"""Given a function, return metadata including a pydantic model representing its
signature.
The use case for this is
```
meta = func_metadata(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.
structured_output: Controls whether the tool's output is structured or unstructured
- If None, auto-detects based on the function's return type annotation
- If True, creates a structured tool (return type annotation permitting)
- If False, unconditionally creates an unstructured tool
If structured, creates a Pydantic model for the function's result based on its annotation.
Supports various return types:
- BaseModel subclasses (used directly)
- Primitive types (str, int, float, bool, bytes, None) - wrapped in a
model with a 'result' field
- TypedDict - converted to a Pydantic model with same fields
- Dataclasses and other annotated classes - converted to Pydantic models
- Generic types (list, dict, Union, etc.) - wrapped in a model with a 'result' field
Returns:
A FuncMetadata object containing:
- arg_model: A pydantic model representing the function's arguments
- output_model: A pydantic model for the return type if output is structured
- output_conversion: Records how function output should be converted before returning.
"""
try:
sig = inspect.signature(func, eval_str=True)
except NameError as e: # pragma: no cover
# This raise could perhaps be skipped, and we (FastMCP) just call
# model_rebuild right before using it 🤷
raise InvalidSignature(f"Unable to evaluate type annotations for callable {func.__name__!r}") from e
params = sig.parameters
dynamic_pydantic_model_params: dict[str, Any] = {}
for param in params.values():
if param.name.startswith("_"): # pragma: no cover
raise InvalidSignature(f"Parameter {param.name} of {func.__name__} cannot start with '_'")
if param.name in skip_names:
continue
annotation = param.annotation if param.annotation is not inspect.Parameter.empty else Any
field_name = param.name
field_kwargs: dict[str, Any] = {}
field_metadata: list[Any] = []
if param.annotation is inspect.Parameter.empty:
field_metadata.append(WithJsonSchema({"title": param.name, "type": "string"}))
# Check if the parameter name conflicts with BaseModel attributes
# This is necessary because Pydantic warns about shadowing parent attributes
if hasattr(BaseModel, field_name) and callable(getattr(BaseModel, field_name)):
# Use an alias to avoid the shadowing warning
field_kwargs["alias"] = field_name
# Use a prefixed field name
field_name = f"field_{field_name}"
if param.default is not inspect.Parameter.empty:
dynamic_pydantic_model_params[field_name] = (
Annotated[(annotation, *field_metadata, Field(**field_kwargs))],
param.default,
)
else:
dynamic_pydantic_model_params[field_name] = Annotated[(annotation, *field_metadata, Field(**field_kwargs))]
arguments_model = create_model(
f"{func.__name__}Arguments",
__base__=ArgModelBase,
**dynamic_pydantic_model_params,
)
if structured_output is False:
return FuncMetadata(arg_model=arguments_model)
# set up structured output support based on return type annotation
if sig.return_annotation is inspect.Parameter.empty and structured_output is True:
raise InvalidSignature(f"Function {func.__name__}: return annotation required for structured output")
try:
inspected_return_ann = inspect_annotation(sig.return_annotation, annotation_source=AnnotationSource.FUNCTION)
except ForbiddenQualifier as e:
raise InvalidSignature(f"Function {func.__name__}: return annotation contains an invalid type qualifier") from e
return_type_expr = inspected_return_ann.type
# `AnnotationSource.FUNCTION` allows no type qualifier to be used, so `return_type_expr` is guaranteed to *not* be
# unknown (i.e. a bare `Final`).
assert return_type_expr is not UNKNOWN
if is_union_origin(get_origin(return_type_expr)):
args = get_args(return_type_expr)
# Check if CallToolResult appears in the union (excluding None for Optional check)
if any(isinstance(arg, type) and issubclass(arg, CallToolResult) for arg in args if arg is not type(None)):
raise InvalidSignature(
f"Function {func.__name__}: CallToolResult cannot be used in Union or Optional types. "
"To return empty results, use: CallToolResult(content=[])"
)
original_annotation: Any
# if the typehint is CallToolResult, the user either intends to return without validation
# or they provided validation as Annotated metadata
if isinstance(return_type_expr, type) and issubclass(return_type_expr, CallToolResult):
if inspected_return_ann.metadata:
return_type_expr = inspected_return_ann.metadata[0]
if len(inspected_return_ann.metadata) >= 2:
# Reconstruct the original annotation, by preserving the remaining metadata,
# i.e. from `Annotated[CallToolResult, ReturnType, Gt(1)]` to
# `Annotated[ReturnType, Gt(1)]`:
original_annotation = Annotated[
(return_type_expr, *inspected_return_ann.metadata[1:])
] # pragma: no cover
else:
# We only had `Annotated[CallToolResult, ReturnType]`, treat the original annotation
# as beging `ReturnType`:
original_annotation = return_type_expr
else:
return FuncMetadata(arg_model=arguments_model)
else:
original_annotation = sig.return_annotation
output_model, output_schema, wrap_output = _try_create_model_and_schema(
original_annotation, return_type_expr, func.__name__
)
if output_model is None and structured_output is True:
# Model creation failed or produced warnings - no structured output
raise InvalidSignature(
f"Function {func.__name__}: return type {return_type_expr} is not serializable for structured output"
)
return FuncMetadata(
arg_model=arguments_model,
output_schema=output_schema,
output_model=output_model,
wrap_output=wrap_output,
)
def _try_create_model_and_schema(
original_annotation: Any,
type_expr: Any,
func_name: str,
) -> tuple[type[BaseModel] | None, dict[str, Any] | None, bool]:
"""Try to create a model and schema for the given annotation without warnings.
Args:
original_annotation: The original return annotation (may be wrapped in `Annotated`).
type_expr: The underlying type expression derived from the return annotation
(`Annotated` and type qualifiers were stripped).
func_name: The name of the function.
Returns:
tuple of (model or None, schema or None, wrap_output)
Model and schema are None if warnings occur or creation fails.
wrap_output is True if the result needs to be wrapped in {"result": ...}
"""
model = None
wrap_output = False
# First handle special case: None
if type_expr is None:
model = _create_wrapped_model(func_name, original_annotation)
wrap_output = True
# Handle GenericAlias types (list[str], dict[str, int], Union[str, int], etc.)
elif isinstance(type_expr, GenericAlias):
origin = get_origin(type_expr)
# Special case: dict with string keys can use RootModel
if origin is dict:
args = get_args(type_expr)
if len(args) == 2 and args[0] is str:
# TODO: should we use the original annotation? We are loosing any potential `Annotated`
# metadata for Pydantic here:
model = _create_dict_model(func_name, type_expr)
else:
# dict with non-str keys needs wrapping
model = _create_wrapped_model(func_name, original_annotation)
wrap_output = True
else:
# All other generic types need wrapping (list, tuple, Union, Optional, etc.)
model = _create_wrapped_model(func_name, original_annotation)
wrap_output = True
# Handle regular type objects
elif isinstance(type_expr, type):
type_annotation = cast(type[Any], type_expr)
# Case 1: BaseModel subclasses (can be used directly)
if issubclass(type_annotation, BaseModel):
model = type_annotation
# Case 2: TypedDicts:
elif is_typeddict(type_annotation):
model = _create_model_from_typeddict(type_annotation)
# Case 3: Primitive types that need wrapping
elif type_annotation in (str, int, float, bool, bytes, type(None)):
model = _create_wrapped_model(func_name, original_annotation)
wrap_output = True
# Case 4: Other class types (dataclasses, regular classes with annotations)
else:
type_hints = get_type_hints(type_annotation)
if type_hints:
# Classes with type hints can be converted to Pydantic models
model = _create_model_from_class(type_annotation, type_hints)
# Classes without type hints are not serializable - model remains None
# Handle any other types not covered above
else:
# This includes typing constructs that aren't GenericAlias in Python 3.10
# (e.g., Union, Optional in some Python versions)
model = _create_wrapped_model(func_name, original_annotation)
wrap_output = True
if model:
# If we successfully created a model, try to get its schema
# Use StrictJsonSchema to raise exceptions instead of warnings
try:
schema = model.model_json_schema(schema_generator=StrictJsonSchema)
except (TypeError, ValueError, pydantic_core.SchemaError, pydantic_core.ValidationError) as e:
# These are expected errors when a type can't be converted to a Pydantic schema
# TypeError: When Pydantic can't handle the type
# ValueError: When there are issues with the type definition (including our custom warnings)
# SchemaError: When Pydantic can't build a schema
# ValidationError: When validation fails
logger.info(f"Cannot create schema for type {type_expr} in {func_name}: {type(e).__name__}: {e}")
return None, None, False
return model, schema, wrap_output
return None, None, False
_no_default = object()
def _create_model_from_class(cls: type[Any], type_hints: dict[str, Any]) -> type[BaseModel]:
"""Create a Pydantic model from an ordinary class.
The created model will:
- Have the same name as the class
- Have fields with the same names and types as the class's fields
- Include all fields whose type does not include None in the set of required fields
Precondition: cls must have type hints (i.e., `type_hints` is non-empty)
"""
model_fields: dict[str, Any] = {}
for field_name, field_type in type_hints.items():
if field_name.startswith("_"): # pragma: no cover
continue
default = getattr(cls, field_name, _no_default)
if default is _no_default:
model_fields[field_name] = field_type
else:
model_fields[field_name] = (field_type, default)
return create_model(cls.__name__, __config__=ConfigDict(from_attributes=True), **model_fields)
def _create_model_from_typeddict(td_type: type[Any]) -> type[BaseModel]:
"""Create a Pydantic model from a TypedDict.
The created model will have the same name and fields as the TypedDict.
"""
type_hints = get_type_hints(td_type)
required_keys = getattr(td_type, "__required_keys__", set(type_hints.keys()))
model_fields: dict[str, Any] = {}
for field_name, field_type in type_hints.items():
if field_name not in required_keys:
# For optional TypedDict fields, set default=None
# This makes them not required in the Pydantic model
# The model should use exclude_unset=True when dumping to get TypedDict semantics
model_fields[field_name] = (field_type, None)
else:
model_fields[field_name] = field_type
return create_model(td_type.__name__, **model_fields)
def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:
"""Create a model that wraps a type in a 'result' field.
This is used for primitive types, generic types like list/dict, etc.
"""
model_name = f"{func_name}Output"
return create_model(model_name, result=annotation)
def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]:
"""Create a RootModel for dict[str, T] types."""
class DictModel(RootModel[dict_annotation]):
pass
# Give it a meaningful name
DictModel.__name__ = f"{func_name}DictOutput"
DictModel.__qualname__ = f"{func_name}DictOutput"
return DictModel
def _convert_to_content(
result: Any,
) -> Sequence[ContentBlock]:
"""
Convert a result to a sequence of content objects.
Note: This conversion logic comes from previous versions of FastMCP and is being
retained for purposes of backwards compatibility. It produces different unstructured
output than the lowlevel server tool call handler, which just serializes structured
content verbatim.
"""
if result is None: # pragma: no cover
return []
if isinstance(result, ContentBlock):
return [result]
if isinstance(result, Image):
return [result.to_image_content()]
if isinstance(result, Audio):
return [result.to_audio_content()]
if isinstance(result, list | tuple):
return list(
chain.from_iterable(
_convert_to_content(item)
for item in result # type: ignore
)
)
if not isinstance(result, str):
result = pydantic_core.to_json(result, fallback=str, indent=2).decode()
return [TextContent(type="text", text=result)]
@@ -0,0 +1,43 @@
"""Logging utilities for FastMCP."""
import logging
from typing import Literal
def get_logger(name: str) -> logging.Logger:
"""Get a logger nested under MCPnamespace.
Args:
name: the name of the logger, which will be prefixed with 'FastMCP.'
Returns:
a configured logger instance
"""
return logging.getLogger(name)
def configure_logging(
level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO",
) -> None:
"""Configure logging for MCP.
Args:
level: the log level to use
"""
handlers: list[logging.Handler] = []
try: # pragma: no cover
from rich.console import Console
from rich.logging import RichHandler
handlers.append(RichHandler(console=Console(stderr=True), rich_tracebacks=True))
except ImportError: # pragma: no cover
pass
if not handlers: # pragma: no cover
handlers.append(logging.StreamHandler())
logging.basicConfig(
level=level,
format="%(message)s",
handlers=handlers,
)
@@ -0,0 +1,101 @@
"""Common types used across FastMCP."""
import base64
from pathlib import Path
from mcp.types import AudioContent, ImageContent
class Image:
"""Helper class for returning images from tools."""
def __init__(
self,
path: str | Path | None = None,
data: bytes | None = None,
format: str | None = None,
):
if path is None and data is None: # pragma: no cover
raise ValueError("Either path or data must be provided")
if path is not None and data is not None: # pragma: no cover
raise ValueError("Only one of path or data can be provided")
self.path = Path(path) if path else None
self.data = data
self._format = format
self._mime_type = self._get_mime_type()
def _get_mime_type(self) -> str:
"""Get MIME type from format or guess from file extension."""
if self._format: # pragma: no cover
return f"image/{self._format.lower()}"
if self.path:
suffix = self.path.suffix.lower()
return {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
}.get(suffix, "application/octet-stream")
return "image/png" # pragma: no cover # default for raw binary data
def to_image_content(self) -> ImageContent:
"""Convert to MCP ImageContent."""
if self.path:
with open(self.path, "rb") as f:
data = base64.b64encode(f.read()).decode()
elif self.data is not None: # pragma: no cover
data = base64.b64encode(self.data).decode()
else: # pragma: no cover
raise ValueError("No image data available")
return ImageContent(type="image", data=data, mimeType=self._mime_type)
class Audio:
"""Helper class for returning audio from tools."""
def __init__(
self,
path: str | Path | None = None,
data: bytes | None = None,
format: str | None = None,
):
if not bool(path) ^ bool(data): # pragma: no cover
raise ValueError("Either path or data can be provided")
self.path = Path(path) if path else None
self.data = data
self._format = format
self._mime_type = self._get_mime_type()
def _get_mime_type(self) -> str:
"""Get MIME type from format or guess from file extension."""
if self._format: # pragma: no cover
return f"audio/{self._format.lower()}"
if self.path:
suffix = self.path.suffix.lower()
return {
".wav": "audio/wav",
".mp3": "audio/mpeg",
".ogg": "audio/ogg",
".flac": "audio/flac",
".aac": "audio/aac",
".m4a": "audio/mp4",
}.get(suffix, "application/octet-stream")
return "audio/wav" # pragma: no cover # default for raw binary data
def to_audio_content(self) -> AudioContent:
"""Convert to MCP AudioContent."""
if self.path:
with open(self.path, "rb") as f:
data = base64.b64encode(f.read()).decode()
elif self.data is not None: # pragma: no cover
data = base64.b64encode(self.data).decode()
else: # pragma: no cover
raise ValueError("No audio data available")
return AudioContent(type="audio", data=data, mimeType=self._mime_type)
@@ -0,0 +1,3 @@
from .server import NotificationOptions, Server
__all__ = ["Server", "NotificationOptions"]
@@ -0,0 +1,288 @@
"""Experimental handlers for the low-level MCP server.
WARNING: These APIs are experimental and may change without notice.
"""
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING
from mcp.server.experimental.task_support import TaskSupport
from mcp.server.lowlevel.func_inspection import create_call_wrapper
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.helpers import cancel_task
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.types import (
INVALID_PARAMS,
CancelTaskRequest,
CancelTaskResult,
ErrorData,
GetTaskPayloadRequest,
GetTaskPayloadResult,
GetTaskRequest,
GetTaskResult,
ListTasksRequest,
ListTasksResult,
ServerCapabilities,
ServerResult,
ServerTasksCapability,
ServerTasksRequestsCapability,
TasksCancelCapability,
TasksListCapability,
TasksToolsCapability,
)
if TYPE_CHECKING:
from mcp.server.lowlevel.server import Server
logger = logging.getLogger(__name__)
class ExperimentalHandlers:
"""Experimental request/notification handlers.
WARNING: These APIs are experimental and may change without notice.
"""
def __init__(
self,
server: Server,
request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]],
notification_handlers: dict[type, Callable[..., Awaitable[None]]],
):
self._server = server
self._request_handlers = request_handlers
self._notification_handlers = notification_handlers
self._task_support: TaskSupport | None = None
@property
def task_support(self) -> TaskSupport | None:
"""Get the task support configuration, if enabled."""
return self._task_support
def update_capabilities(self, capabilities: ServerCapabilities) -> None:
# Only add tasks capability if handlers are registered
if not any(
req_type in self._request_handlers
for req_type in [GetTaskRequest, ListTasksRequest, CancelTaskRequest, GetTaskPayloadRequest]
):
return
capabilities.tasks = ServerTasksCapability()
if ListTasksRequest in self._request_handlers:
capabilities.tasks.list = TasksListCapability()
if CancelTaskRequest in self._request_handlers:
capabilities.tasks.cancel = TasksCancelCapability()
capabilities.tasks.requests = ServerTasksRequestsCapability(
tools=TasksToolsCapability()
) # assuming always supported for now
def enable_tasks(
self,
store: TaskStore | None = None,
queue: TaskMessageQueue | None = None,
) -> TaskSupport:
"""
Enable experimental task support.
This sets up the task infrastructure and auto-registers default handlers
for tasks/get, tasks/result, tasks/list, and tasks/cancel.
Args:
store: Custom TaskStore implementation (defaults to InMemoryTaskStore)
queue: Custom TaskMessageQueue implementation (defaults to InMemoryTaskMessageQueue)
Returns:
The TaskSupport configuration object
Example:
# Simple in-memory setup
server.experimental.enable_tasks()
# Custom store/queue for distributed systems
server.experimental.enable_tasks(
store=RedisTaskStore(redis_url),
queue=RedisTaskMessageQueue(redis_url),
)
WARNING: This API is experimental and may change without notice.
"""
if store is None:
store = InMemoryTaskStore()
if queue is None:
queue = InMemoryTaskMessageQueue()
self._task_support = TaskSupport(store=store, queue=queue)
# Auto-register default handlers
self._register_default_task_handlers()
return self._task_support
def _register_default_task_handlers(self) -> None:
"""Register default handlers for task operations."""
assert self._task_support is not None
support = self._task_support
# Register get_task handler if not already registered
if GetTaskRequest not in self._request_handlers:
async def _default_get_task(req: GetTaskRequest) -> ServerResult:
task = await support.store.get_task(req.params.taskId)
if task is None:
raise McpError(
ErrorData(
code=INVALID_PARAMS,
message=f"Task not found: {req.params.taskId}",
)
)
return ServerResult(
GetTaskResult(
taskId=task.taskId,
status=task.status,
statusMessage=task.statusMessage,
createdAt=task.createdAt,
lastUpdatedAt=task.lastUpdatedAt,
ttl=task.ttl,
pollInterval=task.pollInterval,
)
)
self._request_handlers[GetTaskRequest] = _default_get_task
# Register get_task_result handler if not already registered
if GetTaskPayloadRequest not in self._request_handlers:
async def _default_get_task_result(req: GetTaskPayloadRequest) -> ServerResult:
ctx = self._server.request_context
result = await support.handler.handle(req, ctx.session, ctx.request_id)
return ServerResult(result)
self._request_handlers[GetTaskPayloadRequest] = _default_get_task_result
# Register list_tasks handler if not already registered
if ListTasksRequest not in self._request_handlers:
async def _default_list_tasks(req: ListTasksRequest) -> ServerResult:
cursor = req.params.cursor if req.params else None
tasks, next_cursor = await support.store.list_tasks(cursor)
return ServerResult(ListTasksResult(tasks=tasks, nextCursor=next_cursor))
self._request_handlers[ListTasksRequest] = _default_list_tasks
# Register cancel_task handler if not already registered
if CancelTaskRequest not in self._request_handlers:
async def _default_cancel_task(req: CancelTaskRequest) -> ServerResult:
result = await cancel_task(support.store, req.params.taskId)
return ServerResult(result)
self._request_handlers[CancelTaskRequest] = _default_cancel_task
def list_tasks(
self,
) -> Callable[
[Callable[[ListTasksRequest], Awaitable[ListTasksResult]]],
Callable[[ListTasksRequest], Awaitable[ListTasksResult]],
]:
"""Register a handler for listing tasks.
WARNING: This API is experimental and may change without notice.
"""
def decorator(
func: Callable[[ListTasksRequest], Awaitable[ListTasksResult]],
) -> Callable[[ListTasksRequest], Awaitable[ListTasksResult]]:
logger.debug("Registering handler for ListTasksRequest")
wrapper = create_call_wrapper(func, ListTasksRequest)
async def handler(req: ListTasksRequest) -> ServerResult:
result = await wrapper(req)
return ServerResult(result)
self._request_handlers[ListTasksRequest] = handler
return func
return decorator
def get_task(
self,
) -> Callable[
[Callable[[GetTaskRequest], Awaitable[GetTaskResult]]], Callable[[GetTaskRequest], Awaitable[GetTaskResult]]
]:
"""Register a handler for getting task status.
WARNING: This API is experimental and may change without notice.
"""
def decorator(
func: Callable[[GetTaskRequest], Awaitable[GetTaskResult]],
) -> Callable[[GetTaskRequest], Awaitable[GetTaskResult]]:
logger.debug("Registering handler for GetTaskRequest")
wrapper = create_call_wrapper(func, GetTaskRequest)
async def handler(req: GetTaskRequest) -> ServerResult:
result = await wrapper(req)
return ServerResult(result)
self._request_handlers[GetTaskRequest] = handler
return func
return decorator
def get_task_result(
self,
) -> Callable[
[Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]],
Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
]:
"""Register a handler for getting task results/payload.
WARNING: This API is experimental and may change without notice.
"""
def decorator(
func: Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
) -> Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]:
logger.debug("Registering handler for GetTaskPayloadRequest")
wrapper = create_call_wrapper(func, GetTaskPayloadRequest)
async def handler(req: GetTaskPayloadRequest) -> ServerResult:
result = await wrapper(req)
return ServerResult(result)
self._request_handlers[GetTaskPayloadRequest] = handler
return func
return decorator
def cancel_task(
self,
) -> Callable[
[Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]],
Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
]:
"""Register a handler for cancelling tasks.
WARNING: This API is experimental and may change without notice.
"""
def decorator(
func: Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
) -> Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]:
logger.debug("Registering handler for CancelTaskRequest")
wrapper = create_call_wrapper(func, CancelTaskRequest)
async def handler(req: CancelTaskRequest) -> ServerResult:
result = await wrapper(req)
return ServerResult(result)
self._request_handlers[CancelTaskRequest] = handler
return func
return decorator
@@ -0,0 +1,54 @@
import inspect
from collections.abc import Callable
from typing import Any, TypeVar, get_type_hints
T = TypeVar("T")
R = TypeVar("R")
def create_call_wrapper(func: Callable[..., R], request_type: type[T]) -> Callable[[T], R]:
"""
Create a wrapper function that knows how to call func with the request object.
Returns a wrapper function that takes the request and calls func appropriately.
The wrapper handles three calling patterns:
1. Positional-only parameter typed as request_type (no default): func(req)
2. Positional/keyword parameter typed as request_type (no default): func(**{param_name: req})
3. No request parameter or parameter with default: func()
"""
try:
sig = inspect.signature(func)
type_hints = get_type_hints(func)
except (ValueError, TypeError, NameError): # pragma: no cover
return lambda _: func()
# Check for positional-only parameter typed as request_type
for param_name, param in sig.parameters.items():
if param.kind == inspect.Parameter.POSITIONAL_ONLY:
param_type = type_hints.get(param_name)
if param_type == request_type: # pragma: no branch
# Check if it has a default - if so, treat as old style
if param.default is not inspect.Parameter.empty: # pragma: no cover
return lambda _: func()
# Found positional-only parameter with correct type and no default
return lambda req: func(req)
# Check for any positional/keyword parameter typed as request_type
for param_name, param in sig.parameters.items():
if param.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY): # pragma: no branch
param_type = type_hints.get(param_name)
if param_type == request_type:
# Check if it has a default - if so, treat as old style
if param.default is not inspect.Parameter.empty: # pragma: no cover
return lambda _: func()
# Found keyword parameter with correct type and no default
# Need to capture param_name in closure properly
def make_keyword_wrapper(name: str) -> Callable[[Any], Any]:
return lambda req: func(**{name: req})
return make_keyword_wrapper(param_name)
# No request parameter found - use old style
return lambda _: func()
@@ -0,0 +1,11 @@
from dataclasses import dataclass
from typing import Any
@dataclass
class ReadResourceContents:
"""Contents returned from a read_resource call."""
content: str | bytes
mime_type: str | None = None
meta: dict[str, Any] | None = None
@@ -0,0 +1,803 @@
"""
MCP Server Module
This module provides a framework for creating an MCP (Model Context Protocol) server.
It allows you to easily define and handle various types of requests and notifications
in an asynchronous manner.
Usage:
1. Create a Server instance:
server = Server("your_server_name")
2. Define request handlers using decorators:
@server.list_prompts()
async def handle_list_prompts(request: types.ListPromptsRequest) -> types.ListPromptsResult:
# Implementation
@server.get_prompt()
async def handle_get_prompt(
name: str, arguments: dict[str, str] | None
) -> types.GetPromptResult:
# Implementation
@server.list_tools()
async def handle_list_tools(request: types.ListToolsRequest) -> types.ListToolsResult:
# Implementation
@server.call_tool()
async def handle_call_tool(
name: str, arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
# Implementation
@server.list_resource_templates()
async def handle_list_resource_templates() -> list[types.ResourceTemplate]:
# Implementation
3. Define notification handlers if needed:
@server.progress_notification()
async def handle_progress(
progress_token: str | int, progress: float, total: float | None,
message: str | None
) -> None:
# Implementation
4. Run the server:
async def main():
async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
InitializationOptions(
server_name="your_server_name",
server_version="your_version",
capabilities=server.get_capabilities(
notification_options=NotificationOptions(),
experimental_capabilities={},
),
),
)
asyncio.run(main())
The Server class provides methods to register handlers for various MCP requests and
notifications. It automatically manages the request context and handles incoming
messages from the client.
"""
from __future__ import annotations as _annotations
import base64
import contextvars
import json
import logging
import warnings
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
from typing import Any, Generic, TypeAlias, cast
import anyio
import jsonschema
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import AnyUrl
from typing_extensions import TypeVar
import mcp.types as types
from mcp.server.experimental.request_context import Experimental
from mcp.server.lowlevel.experimental import ExperimentalHandlers
from mcp.server.lowlevel.func_inspection import create_call_wrapper
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.context import RequestContext
from mcp.shared.exceptions import McpError, UrlElicitationRequiredError
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.shared.session import RequestResponder
from mcp.shared.tool_name_validation import validate_and_warn_tool_name
logger = logging.getLogger(__name__)
LifespanResultT = TypeVar("LifespanResultT", default=Any)
RequestT = TypeVar("RequestT", default=Any)
# type aliases for tool call results
StructuredContent: TypeAlias = dict[str, Any]
UnstructuredContent: TypeAlias = Iterable[types.ContentBlock]
CombinationContent: TypeAlias = tuple[UnstructuredContent, StructuredContent]
# This will be properly typed in each Server instance's context
request_ctx: contextvars.ContextVar[RequestContext[ServerSession, Any, Any]] = contextvars.ContextVar("request_ctx")
class NotificationOptions:
def __init__(
self,
prompts_changed: bool = False,
resources_changed: bool = False,
tools_changed: bool = False,
):
self.prompts_changed = prompts_changed
self.resources_changed = resources_changed
self.tools_changed = tools_changed
@asynccontextmanager
async def lifespan(_: Server[LifespanResultT, RequestT]) -> AsyncIterator[dict[str, Any]]:
"""Default lifespan context manager that does nothing.
Args:
server: The server instance this lifespan is managing
Returns:
An empty context object
"""
yield {}
class Server(Generic[LifespanResultT, RequestT]):
def __init__(
self,
name: str,
version: str | None = None,
instructions: str | None = None,
website_url: str | None = None,
icons: list[types.Icon] | None = None,
lifespan: Callable[
[Server[LifespanResultT, RequestT]],
AbstractAsyncContextManager[LifespanResultT],
] = lifespan,
):
self.name = name
self.version = version
self.instructions = instructions
self.website_url = website_url
self.icons = icons
self.lifespan = lifespan
self.request_handlers: dict[type, Callable[..., Awaitable[types.ServerResult]]] = {
types.PingRequest: _ping_handler,
}
self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
self._tool_cache: dict[str, types.Tool] = {}
self._experimental_handlers: ExperimentalHandlers | None = None
logger.debug("Initializing server %r", name)
def create_initialization_options(
self,
notification_options: NotificationOptions | None = None,
experimental_capabilities: dict[str, dict[str, Any]] | None = None,
) -> InitializationOptions:
"""Create initialization options from this server instance."""
def pkg_version(package: str) -> str:
try:
from importlib.metadata import version
return version(package)
except Exception: # pragma: no cover
pass
return "unknown" # pragma: no cover
return InitializationOptions(
server_name=self.name,
server_version=self.version if self.version else pkg_version("mcp"),
capabilities=self.get_capabilities(
notification_options or NotificationOptions(),
experimental_capabilities or {},
),
instructions=self.instructions,
website_url=self.website_url,
icons=self.icons,
)
def get_capabilities(
self,
notification_options: NotificationOptions,
experimental_capabilities: dict[str, dict[str, Any]],
) -> types.ServerCapabilities:
"""Convert existing handlers to a ServerCapabilities object."""
prompts_capability = None
resources_capability = None
tools_capability = None
logging_capability = None
completions_capability = None
# Set prompt capabilities if handler exists
if types.ListPromptsRequest in self.request_handlers:
prompts_capability = types.PromptsCapability(listChanged=notification_options.prompts_changed)
# Set resource capabilities if handler exists
if types.ListResourcesRequest in self.request_handlers:
resources_capability = types.ResourcesCapability(
subscribe=False, listChanged=notification_options.resources_changed
)
# Set tool capabilities if handler exists
if types.ListToolsRequest in self.request_handlers:
tools_capability = types.ToolsCapability(listChanged=notification_options.tools_changed)
# Set logging capabilities if handler exists
if types.SetLevelRequest in self.request_handlers: # pragma: no cover
logging_capability = types.LoggingCapability()
# Set completions capabilities if handler exists
if types.CompleteRequest in self.request_handlers:
completions_capability = types.CompletionsCapability()
capabilities = types.ServerCapabilities(
prompts=prompts_capability,
resources=resources_capability,
tools=tools_capability,
logging=logging_capability,
experimental=experimental_capabilities,
completions=completions_capability,
)
if self._experimental_handlers:
self._experimental_handlers.update_capabilities(capabilities)
return capabilities
@property
def request_context(
self,
) -> RequestContext[ServerSession, LifespanResultT, RequestT]:
"""If called outside of a request context, this will raise a LookupError."""
return request_ctx.get()
@property
def experimental(self) -> ExperimentalHandlers:
"""Experimental APIs for tasks and other features.
WARNING: These APIs are experimental and may change without notice.
"""
# We create this inline so we only add these capabilities _if_ they're actually used
if self._experimental_handlers is None:
self._experimental_handlers = ExperimentalHandlers(self, self.request_handlers, self.notification_handlers)
return self._experimental_handlers
def list_prompts(self):
def decorator(
func: Callable[[], Awaitable[list[types.Prompt]]]
| Callable[[types.ListPromptsRequest], Awaitable[types.ListPromptsResult]],
):
logger.debug("Registering handler for PromptListRequest")
wrapper = create_call_wrapper(func, types.ListPromptsRequest)
async def handler(req: types.ListPromptsRequest):
result = await wrapper(req)
# Handle both old style (list[Prompt]) and new style (ListPromptsResult)
if isinstance(result, types.ListPromptsResult):
return types.ServerResult(result)
else:
# Old style returns list[Prompt]
return types.ServerResult(types.ListPromptsResult(prompts=result))
self.request_handlers[types.ListPromptsRequest] = handler
return func
return decorator
def get_prompt(self):
def decorator(
func: Callable[[str, dict[str, str] | None], Awaitable[types.GetPromptResult]],
):
logger.debug("Registering handler for GetPromptRequest")
async def handler(req: types.GetPromptRequest):
prompt_get = await func(req.params.name, req.params.arguments)
return types.ServerResult(prompt_get)
self.request_handlers[types.GetPromptRequest] = handler
return func
return decorator
def list_resources(self):
def decorator(
func: Callable[[], Awaitable[list[types.Resource]]]
| Callable[[types.ListResourcesRequest], Awaitable[types.ListResourcesResult]],
):
logger.debug("Registering handler for ListResourcesRequest")
wrapper = create_call_wrapper(func, types.ListResourcesRequest)
async def handler(req: types.ListResourcesRequest):
result = await wrapper(req)
# Handle both old style (list[Resource]) and new style (ListResourcesResult)
if isinstance(result, types.ListResourcesResult):
return types.ServerResult(result)
else:
# Old style returns list[Resource]
return types.ServerResult(types.ListResourcesResult(resources=result))
self.request_handlers[types.ListResourcesRequest] = handler
return func
return decorator
def list_resource_templates(self):
def decorator(func: Callable[[], Awaitable[list[types.ResourceTemplate]]]):
logger.debug("Registering handler for ListResourceTemplatesRequest")
async def handler(_: Any):
templates = await func()
return types.ServerResult(types.ListResourceTemplatesResult(resourceTemplates=templates))
self.request_handlers[types.ListResourceTemplatesRequest] = handler
return func
return decorator
def read_resource(self):
def decorator(
func: Callable[[AnyUrl], Awaitable[str | bytes | Iterable[ReadResourceContents]]],
):
logger.debug("Registering handler for ReadResourceRequest")
async def handler(req: types.ReadResourceRequest):
result = await func(req.params.uri)
def create_content(data: str | bytes, mime_type: str | None, meta: dict[str, Any] | None = None):
# Note: ResourceContents uses Field(alias="_meta"), so we must use the alias key
meta_kwargs: dict[str, Any] = {"_meta": meta} if meta is not None else {}
match data:
case str() as data:
return types.TextResourceContents(
uri=req.params.uri,
text=data,
mimeType=mime_type or "text/plain",
**meta_kwargs,
)
case bytes() as data: # pragma: no cover
return types.BlobResourceContents(
uri=req.params.uri,
blob=base64.b64encode(data).decode(),
mimeType=mime_type or "application/octet-stream",
**meta_kwargs,
)
match result:
case str() | bytes() as data: # pragma: no cover
warnings.warn(
"Returning str or bytes from read_resource is deprecated. "
"Use Iterable[ReadResourceContents] instead.",
DeprecationWarning,
stacklevel=2,
)
content = create_content(data, None)
case Iterable() as contents:
contents_list = [
create_content(
content_item.content, content_item.mime_type, getattr(content_item, "meta", None)
)
for content_item in contents
]
return types.ServerResult(
types.ReadResourceResult(
contents=contents_list,
)
)
case _: # pragma: no cover
raise ValueError(f"Unexpected return type from read_resource: {type(result)}")
return types.ServerResult( # pragma: no cover
types.ReadResourceResult(
contents=[content],
)
)
self.request_handlers[types.ReadResourceRequest] = handler
return func
return decorator
def set_logging_level(self): # pragma: no cover
def decorator(func: Callable[[types.LoggingLevel], Awaitable[None]]):
logger.debug("Registering handler for SetLevelRequest")
async def handler(req: types.SetLevelRequest):
await func(req.params.level)
return types.ServerResult(types.EmptyResult())
self.request_handlers[types.SetLevelRequest] = handler
return func
return decorator
def subscribe_resource(self): # pragma: no cover
def decorator(func: Callable[[AnyUrl], Awaitable[None]]):
logger.debug("Registering handler for SubscribeRequest")
async def handler(req: types.SubscribeRequest):
await func(req.params.uri)
return types.ServerResult(types.EmptyResult())
self.request_handlers[types.SubscribeRequest] = handler
return func
return decorator
def unsubscribe_resource(self): # pragma: no cover
def decorator(func: Callable[[AnyUrl], Awaitable[None]]):
logger.debug("Registering handler for UnsubscribeRequest")
async def handler(req: types.UnsubscribeRequest):
await func(req.params.uri)
return types.ServerResult(types.EmptyResult())
self.request_handlers[types.UnsubscribeRequest] = handler
return func
return decorator
def list_tools(self):
def decorator(
func: Callable[[], Awaitable[list[types.Tool]]]
| Callable[[types.ListToolsRequest], Awaitable[types.ListToolsResult]],
):
logger.debug("Registering handler for ListToolsRequest")
wrapper = create_call_wrapper(func, types.ListToolsRequest)
async def handler(req: types.ListToolsRequest):
result = await wrapper(req)
# Handle both old style (list[Tool]) and new style (ListToolsResult)
if isinstance(result, types.ListToolsResult): # pragma: no cover
# Refresh the tool cache with returned tools
for tool in result.tools:
validate_and_warn_tool_name(tool.name)
self._tool_cache[tool.name] = tool
return types.ServerResult(result)
else:
# Old style returns list[Tool]
# Clear and refresh the entire tool cache
self._tool_cache.clear()
for tool in result:
validate_and_warn_tool_name(tool.name)
self._tool_cache[tool.name] = tool
return types.ServerResult(types.ListToolsResult(tools=result))
self.request_handlers[types.ListToolsRequest] = handler
return func
return decorator
def _make_error_result(self, error_message: str) -> types.ServerResult:
"""Create a ServerResult with an error CallToolResult."""
return types.ServerResult(
types.CallToolResult(
content=[types.TextContent(type="text", text=error_message)],
isError=True,
)
)
async def _get_cached_tool_definition(self, tool_name: str) -> types.Tool | None:
"""Get tool definition from cache, refreshing if necessary.
Returns the Tool object if found, None otherwise.
"""
if tool_name not in self._tool_cache:
if types.ListToolsRequest in self.request_handlers:
logger.debug("Tool cache miss for %s, refreshing cache", tool_name)
await self.request_handlers[types.ListToolsRequest](None)
tool = self._tool_cache.get(tool_name)
if tool is None:
logger.warning("Tool '%s' not listed, no validation will be performed", tool_name)
return tool
def call_tool(self, *, validate_input: bool = True):
"""Register a tool call handler.
Args:
validate_input: If True, validates input against inputSchema. Default is True.
The handler validates input against inputSchema (if validate_input=True), calls the tool function,
and builds a CallToolResult with the results:
- Unstructured content (iterable of ContentBlock): returned in content
- Structured content (dict): returned in structuredContent, serialized JSON text returned in content
- Both: returned in content and structuredContent
If outputSchema is defined, validates structuredContent or errors if missing.
"""
def decorator(
func: Callable[
...,
Awaitable[
UnstructuredContent
| StructuredContent
| CombinationContent
| types.CallToolResult
| types.CreateTaskResult
],
],
):
logger.debug("Registering handler for CallToolRequest")
async def handler(req: types.CallToolRequest):
try:
tool_name = req.params.name
arguments = req.params.arguments or {}
tool = await self._get_cached_tool_definition(tool_name)
# input validation
if validate_input and tool:
try:
jsonschema.validate(instance=arguments, schema=tool.inputSchema)
except jsonschema.ValidationError as e:
return self._make_error_result(f"Input validation error: {e.message}")
# tool call
results = await func(tool_name, arguments)
# output normalization
unstructured_content: UnstructuredContent
maybe_structured_content: StructuredContent | None
if isinstance(results, types.CallToolResult):
return types.ServerResult(results)
elif isinstance(results, types.CreateTaskResult):
# Task-augmented execution returns task info instead of result
return types.ServerResult(results)
elif isinstance(results, tuple) and len(results) == 2:
# tool returned both structured and unstructured content
unstructured_content, maybe_structured_content = cast(CombinationContent, results)
elif isinstance(results, dict):
# tool returned structured content only
maybe_structured_content = cast(StructuredContent, results)
unstructured_content = [types.TextContent(type="text", text=json.dumps(results, indent=2))]
elif hasattr(results, "__iter__"): # pragma: no cover
# tool returned unstructured content only
unstructured_content = cast(UnstructuredContent, results)
maybe_structured_content = None
else: # pragma: no cover
return self._make_error_result(f"Unexpected return type from tool: {type(results).__name__}")
# output validation
if tool and tool.outputSchema is not None:
if maybe_structured_content is None:
return self._make_error_result(
"Output validation error: outputSchema defined but no structured output returned"
)
else:
try:
jsonschema.validate(instance=maybe_structured_content, schema=tool.outputSchema)
except jsonschema.ValidationError as e:
return self._make_error_result(f"Output validation error: {e.message}")
# result
return types.ServerResult(
types.CallToolResult(
content=list(unstructured_content),
structuredContent=maybe_structured_content,
isError=False,
)
)
except UrlElicitationRequiredError:
# Re-raise UrlElicitationRequiredError so it can be properly handled
# by _handle_request, which converts it to an error response with code -32042
raise
except Exception as e:
return self._make_error_result(str(e))
self.request_handlers[types.CallToolRequest] = handler
return func
return decorator
def progress_notification(self):
def decorator(
func: Callable[[str | int, float, float | None, str | None], Awaitable[None]],
):
logger.debug("Registering handler for ProgressNotification")
async def handler(req: types.ProgressNotification):
await func(
req.params.progressToken,
req.params.progress,
req.params.total,
req.params.message,
)
self.notification_handlers[types.ProgressNotification] = handler
return func
return decorator
def completion(self):
"""Provides completions for prompts and resource templates"""
def decorator(
func: Callable[
[
types.PromptReference | types.ResourceTemplateReference,
types.CompletionArgument,
types.CompletionContext | None,
],
Awaitable[types.Completion | None],
],
):
logger.debug("Registering handler for CompleteRequest")
async def handler(req: types.CompleteRequest):
completion = await func(req.params.ref, req.params.argument, req.params.context)
return types.ServerResult(
types.CompleteResult(
completion=completion
if completion is not None
else types.Completion(values=[], total=None, hasMore=None),
)
)
self.request_handlers[types.CompleteRequest] = handler
return func
return decorator
async def run(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
initialization_options: InitializationOptions,
# When False, exceptions are returned as messages to the client.
# When True, exceptions are raised, which will cause the server to shut down
# but also make tracing exceptions much easier during testing and when using
# in-process servers.
raise_exceptions: bool = False,
# When True, the server is stateless and
# clients can perform initialization with any node. The client must still follow
# the initialization lifecycle, but can do so with any available node
# rather than requiring initialization for each connection.
stateless: bool = False,
):
async with AsyncExitStack() as stack:
lifespan_context = await stack.enter_async_context(self.lifespan(self))
session = await stack.enter_async_context(
ServerSession(
read_stream,
write_stream,
initialization_options,
stateless=stateless,
)
)
# Configure task support for this session if enabled
task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
if task_support is not None:
task_support.configure_session(session)
await stack.enter_async_context(task_support.run())
async with anyio.create_task_group() as tg:
async for message in session.incoming_messages:
logger.debug("Received message: %s", message)
tg.start_soon(
self._handle_message,
message,
session,
lifespan_context,
raise_exceptions,
)
async def _handle_message(
self,
message: RequestResponder[types.ClientRequest, types.ServerResult] | types.ClientNotification | Exception,
session: ServerSession,
lifespan_context: LifespanResultT,
raise_exceptions: bool = False,
):
with warnings.catch_warnings(record=True) as w:
match message:
case RequestResponder(request=types.ClientRequest(root=req)) as responder:
with responder:
await self._handle_request(message, req, session, lifespan_context, raise_exceptions)
case types.ClientNotification(root=notify):
await self._handle_notification(notify)
case Exception(): # pragma: no cover
logger.error(f"Received exception from stream: {message}")
await session.send_log_message(
level="error",
data="Internal Server Error",
logger="mcp.server.exception_handler",
)
if raise_exceptions:
raise message
for warning in w: # pragma: no cover
logger.info("Warning: %s: %s", warning.category.__name__, warning.message)
async def _handle_request(
self,
message: RequestResponder[types.ClientRequest, types.ServerResult],
req: types.ClientRequestType,
session: ServerSession,
lifespan_context: LifespanResultT,
raise_exceptions: bool,
):
logger.info("Processing request of type %s", type(req).__name__)
if handler := self.request_handlers.get(type(req)):
logger.debug("Dispatching request of type %s", type(req).__name__)
token = None
try:
# Extract request context and close_sse_stream from message metadata
request_data = None
close_sse_stream_cb = None
close_standalone_sse_stream_cb = None
if message.message_metadata is not None and isinstance(
message.message_metadata, ServerMessageMetadata
): # pragma: no cover
request_data = message.message_metadata.request_context
close_sse_stream_cb = message.message_metadata.close_sse_stream
close_standalone_sse_stream_cb = message.message_metadata.close_standalone_sse_stream
# Set our global state that can be retrieved via
# app.get_request_context()
client_capabilities = session.client_params.capabilities if session.client_params else None
task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
# Get task metadata from request params if present
task_metadata = None
if hasattr(req, "params") and req.params is not None:
task_metadata = getattr(req.params, "task", None)
token = request_ctx.set(
RequestContext(
message.request_id,
message.request_meta,
session,
lifespan_context,
Experimental(
task_metadata=task_metadata,
_client_capabilities=client_capabilities,
_session=session,
_task_support=task_support,
),
request=request_data,
close_sse_stream=close_sse_stream_cb,
close_standalone_sse_stream=close_standalone_sse_stream_cb,
)
)
response = await handler(req)
except McpError as err: # pragma: no cover
response = err.error
except anyio.get_cancelled_exc_class(): # pragma: no cover
logger.info(
"Request %s cancelled - duplicate response suppressed",
message.request_id,
)
return
except Exception as err: # pragma: no cover
if raise_exceptions:
raise err
response = types.ErrorData(code=0, message=str(err), data=None)
finally:
# Reset the global state after we are done
if token is not None: # pragma: no branch
request_ctx.reset(token)
await message.respond(response)
else: # pragma: no cover
await message.respond(
types.ErrorData(
code=types.METHOD_NOT_FOUND,
message="Method not found",
)
)
logger.debug("Response sent")
async def _handle_notification(self, notify: Any):
if handler := self.notification_handlers.get(type(notify)): # type: ignore
logger.debug("Dispatching notification of type %s", type(notify).__name__)
try:
await handler(notify)
except Exception: # pragma: no cover
logger.exception("Uncaught exception in notification handler")
async def _ping_handler(request: types.PingRequest) -> types.ServerResult:
return types.ServerResult(types.EmptyResult())
@@ -0,0 +1,20 @@
"""
This module provides simpler types to use with the server for managing prompts
and tools.
"""
from pydantic import BaseModel
from mcp.types import (
Icon,
ServerCapabilities,
)
class InitializationOptions(BaseModel):
server_name: str
server_version: str
capabilities: ServerCapabilities
instructions: str | None = None
website_url: str | None = None
icons: list[Icon] | None = None
@@ -0,0 +1,691 @@
"""
ServerSession Module
This module provides the ServerSession class, which manages communication between the
server and client in the MCP (Model Context Protocol) framework. It is most commonly
used in MCP servers to interact with the client.
Common usage pattern:
```
server = Server(name)
@server.call_tool()
async def handle_tool_call(ctx: RequestContext, arguments: dict[str, Any]) -> Any:
# Check client capabilities before proceeding
if ctx.session.check_client_capability(
types.ClientCapabilities(experimental={"advanced_tools": dict()})
):
# Perform advanced tool operations
result = await perform_advanced_tool_operation(arguments)
else:
# Fall back to basic tool operations
result = await perform_basic_tool_operation(arguments)
return result
@server.list_prompts()
async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
# Access session for any necessary checks or operations
if ctx.session.client_params:
# Customize prompts based on client initialization parameters
return generate_custom_prompts(ctx.session.client_params)
else:
return default_prompts
```
The ServerSession class is typically used internally by the Server class and should not
be instantiated directly by users of the MCP framework.
"""
from enum import Enum
from typing import Any, TypeVar, overload
import anyio
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import AnyUrl
import mcp.types as types
from mcp.server.experimental.session_features import ExperimentalServerSessionFeatures
from mcp.server.models import InitializationOptions
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
from mcp.shared.experimental.tasks.capabilities import check_tasks_capability
from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.shared.session import (
BaseSession,
RequestResponder,
)
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
class InitializationState(Enum):
NotInitialized = 1
Initializing = 2
Initialized = 3
ServerSessionT = TypeVar("ServerSessionT", bound="ServerSession")
ServerRequestResponder = (
RequestResponder[types.ClientRequest, types.ServerResult] | types.ClientNotification | Exception
)
class ServerSession(
BaseSession[
types.ServerRequest,
types.ServerNotification,
types.ServerResult,
types.ClientRequest,
types.ClientNotification,
]
):
_initialized: InitializationState = InitializationState.NotInitialized
_client_params: types.InitializeRequestParams | None = None
_experimental_features: ExperimentalServerSessionFeatures | None = None
def __init__(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
init_options: InitializationOptions,
stateless: bool = False,
) -> None:
super().__init__(read_stream, write_stream, types.ClientRequest, types.ClientNotification)
self._initialization_state = (
InitializationState.Initialized if stateless else InitializationState.NotInitialized
)
self._init_options = init_options
self._incoming_message_stream_writer, self._incoming_message_stream_reader = anyio.create_memory_object_stream[
ServerRequestResponder
](0)
self._exit_stack.push_async_callback(lambda: self._incoming_message_stream_reader.aclose())
@property
def client_params(self) -> types.InitializeRequestParams | None:
return self._client_params # pragma: no cover
@property
def experimental(self) -> ExperimentalServerSessionFeatures:
"""Experimental APIs for server→client task operations.
WARNING: These APIs are experimental and may change without notice.
"""
if self._experimental_features is None:
self._experimental_features = ExperimentalServerSessionFeatures(self)
return self._experimental_features
def check_client_capability(self, capability: types.ClientCapabilities) -> bool: # pragma: no cover
"""Check if the client supports a specific capability."""
if self._client_params is None:
return False
client_caps = self._client_params.capabilities
if capability.roots is not None:
if client_caps.roots is None:
return False
if capability.roots.listChanged and not client_caps.roots.listChanged:
return False
if capability.sampling is not None:
if client_caps.sampling is None:
return False
if capability.sampling.context is not None and client_caps.sampling.context is None:
return False
if capability.sampling.tools is not None and client_caps.sampling.tools is None:
return False
if capability.elicitation is not None and client_caps.elicitation is None:
return False
if capability.experimental is not None:
if client_caps.experimental is None:
return False
for exp_key, exp_value in capability.experimental.items():
if exp_key not in client_caps.experimental or client_caps.experimental[exp_key] != exp_value:
return False
if capability.tasks is not None:
if client_caps.tasks is None:
return False
if not check_tasks_capability(capability.tasks, client_caps.tasks):
return False
return True
async def _receive_loop(self) -> None:
async with self._incoming_message_stream_writer:
await super()._receive_loop()
async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]):
match responder.request.root:
case types.InitializeRequest(params=params):
requested_version = params.protocolVersion
self._initialization_state = InitializationState.Initializing
self._client_params = params
with responder:
await responder.respond(
types.ServerResult(
types.InitializeResult(
protocolVersion=requested_version
if requested_version in SUPPORTED_PROTOCOL_VERSIONS
else types.LATEST_PROTOCOL_VERSION,
capabilities=self._init_options.capabilities,
serverInfo=types.Implementation(
name=self._init_options.server_name,
version=self._init_options.server_version,
websiteUrl=self._init_options.website_url,
icons=self._init_options.icons,
),
instructions=self._init_options.instructions,
)
)
)
self._initialization_state = InitializationState.Initialized
case types.PingRequest():
# Ping requests are allowed at any time
pass
case _:
if self._initialization_state != InitializationState.Initialized:
raise RuntimeError("Received request before initialization was complete")
async def _received_notification(self, notification: types.ClientNotification) -> None:
# Need this to avoid ASYNC910
await anyio.lowlevel.checkpoint()
match notification.root:
case types.InitializedNotification():
self._initialization_state = InitializationState.Initialized
case _:
if self._initialization_state != InitializationState.Initialized: # pragma: no cover
raise RuntimeError("Received notification before initialization was complete")
async def send_log_message(
self,
level: types.LoggingLevel,
data: Any,
logger: str | None = None,
related_request_id: types.RequestId | None = None,
) -> None:
"""Send a log message notification."""
await self.send_notification(
types.ServerNotification(
types.LoggingMessageNotification(
params=types.LoggingMessageNotificationParams(
level=level,
data=data,
logger=logger,
),
)
),
related_request_id,
)
async def send_resource_updated(self, uri: AnyUrl) -> None: # pragma: no cover
"""Send a resource updated notification."""
await self.send_notification(
types.ServerNotification(
types.ResourceUpdatedNotification(
params=types.ResourceUpdatedNotificationParams(uri=uri),
)
)
)
@overload
async def create_message(
self,
messages: list[types.SamplingMessage],
*,
max_tokens: int,
system_prompt: str | None = None,
include_context: types.IncludeContext | None = None,
temperature: float | None = None,
stop_sequences: list[str] | None = None,
metadata: dict[str, Any] | None = None,
model_preferences: types.ModelPreferences | None = None,
tools: None = None,
tool_choice: types.ToolChoice | None = None,
related_request_id: types.RequestId | None = None,
) -> types.CreateMessageResult:
"""Overload: Without tools, returns single content."""
...
@overload
async def create_message(
self,
messages: list[types.SamplingMessage],
*,
max_tokens: int,
system_prompt: str | None = None,
include_context: types.IncludeContext | None = None,
temperature: float | None = None,
stop_sequences: list[str] | None = None,
metadata: dict[str, Any] | None = None,
model_preferences: types.ModelPreferences | None = None,
tools: list[types.Tool],
tool_choice: types.ToolChoice | None = None,
related_request_id: types.RequestId | None = None,
) -> types.CreateMessageResultWithTools:
"""Overload: With tools, returns array-capable content."""
...
async def create_message(
self,
messages: list[types.SamplingMessage],
*,
max_tokens: int,
system_prompt: str | None = None,
include_context: types.IncludeContext | None = None,
temperature: float | None = None,
stop_sequences: list[str] | None = None,
metadata: dict[str, Any] | None = None,
model_preferences: types.ModelPreferences | None = None,
tools: list[types.Tool] | None = None,
tool_choice: types.ToolChoice | None = None,
related_request_id: types.RequestId | None = None,
) -> types.CreateMessageResult | types.CreateMessageResultWithTools:
"""Send a sampling/create_message request.
Args:
messages: The conversation messages to send.
max_tokens: Maximum number of tokens to generate.
system_prompt: Optional system prompt.
include_context: Optional context inclusion setting.
Should only be set to "thisServer" or "allServers"
if the client has sampling.context capability.
temperature: Optional sampling temperature.
stop_sequences: Optional stop sequences.
metadata: Optional metadata to pass through to the LLM provider.
model_preferences: Optional model selection preferences.
tools: Optional list of tools the LLM can use during sampling.
Requires client to have sampling.tools capability.
tool_choice: Optional control over tool usage behavior.
Requires client to have sampling.tools capability.
related_request_id: Optional ID of a related request.
Returns:
The sampling result from the client.
Raises:
McpError: If tools are provided but client doesn't support them.
ValueError: If tool_use or tool_result message structure is invalid.
"""
client_caps = self._client_params.capabilities if self._client_params else None
validate_sampling_tools(client_caps, tools, tool_choice)
validate_tool_use_result_messages(messages)
request = types.ServerRequest(
types.CreateMessageRequest(
params=types.CreateMessageRequestParams(
messages=messages,
systemPrompt=system_prompt,
includeContext=include_context,
temperature=temperature,
maxTokens=max_tokens,
stopSequences=stop_sequences,
metadata=metadata,
modelPreferences=model_preferences,
tools=tools,
toolChoice=tool_choice,
),
)
)
metadata_obj = ServerMessageMetadata(related_request_id=related_request_id)
# Use different result types based on whether tools are provided
if tools is not None:
return await self.send_request(
request=request,
result_type=types.CreateMessageResultWithTools,
metadata=metadata_obj,
)
return await self.send_request(
request=request,
result_type=types.CreateMessageResult,
metadata=metadata_obj,
)
async def list_roots(self) -> types.ListRootsResult:
"""Send a roots/list request."""
return await self.send_request(
types.ServerRequest(types.ListRootsRequest()),
types.ListRootsResult,
)
async def elicit(
self,
message: str,
requestedSchema: types.ElicitRequestedSchema,
related_request_id: types.RequestId | None = None,
) -> types.ElicitResult:
"""Send a form mode elicitation/create request.
Args:
message: The message to present to the user
requestedSchema: Schema defining the expected response structure
related_request_id: Optional ID of the request that triggered this elicitation
Returns:
The client's response
Note:
This method is deprecated in favor of elicit_form(). It remains for
backward compatibility but new code should use elicit_form().
"""
return await self.elicit_form(message, requestedSchema, related_request_id)
async def elicit_form(
self,
message: str,
requestedSchema: types.ElicitRequestedSchema,
related_request_id: types.RequestId | None = None,
) -> types.ElicitResult:
"""Send a form mode elicitation/create request.
Args:
message: The message to present to the user
requestedSchema: Schema defining the expected response structure
related_request_id: Optional ID of the request that triggered this elicitation
Returns:
The client's response with form data
"""
return await self.send_request(
types.ServerRequest(
types.ElicitRequest(
params=types.ElicitRequestFormParams(
message=message,
requestedSchema=requestedSchema,
),
)
),
types.ElicitResult,
metadata=ServerMessageMetadata(related_request_id=related_request_id),
)
async def elicit_url(
self,
message: str,
url: str,
elicitation_id: str,
related_request_id: types.RequestId | None = None,
) -> types.ElicitResult:
"""Send a URL mode elicitation/create request.
This directs the user to an external URL for out-of-band interactions
like OAuth flows, credential collection, or payment processing.
Args:
message: Human-readable explanation of why the interaction is needed
url: The URL the user should navigate to
elicitation_id: Unique identifier for tracking this elicitation
related_request_id: Optional ID of the request that triggered this elicitation
Returns:
The client's response indicating acceptance, decline, or cancellation
"""
return await self.send_request(
types.ServerRequest(
types.ElicitRequest(
params=types.ElicitRequestURLParams(
message=message,
url=url,
elicitationId=elicitation_id,
),
)
),
types.ElicitResult,
metadata=ServerMessageMetadata(related_request_id=related_request_id),
)
async def send_ping(self) -> types.EmptyResult: # pragma: no cover
"""Send a ping request."""
return await self.send_request(
types.ServerRequest(types.PingRequest()),
types.EmptyResult,
)
async def send_progress_notification(
self,
progress_token: str | int,
progress: float,
total: float | None = None,
message: str | None = None,
related_request_id: str | None = None,
) -> None:
"""Send a progress notification."""
await self.send_notification(
types.ServerNotification(
types.ProgressNotification(
params=types.ProgressNotificationParams(
progressToken=progress_token,
progress=progress,
total=total,
message=message,
),
)
),
related_request_id,
)
async def send_resource_list_changed(self) -> None: # pragma: no cover
"""Send a resource list changed notification."""
await self.send_notification(types.ServerNotification(types.ResourceListChangedNotification()))
async def send_tool_list_changed(self) -> None: # pragma: no cover
"""Send a tool list changed notification."""
await self.send_notification(types.ServerNotification(types.ToolListChangedNotification()))
async def send_prompt_list_changed(self) -> None: # pragma: no cover
"""Send a prompt list changed notification."""
await self.send_notification(types.ServerNotification(types.PromptListChangedNotification()))
async def send_elicit_complete(
self,
elicitation_id: str,
related_request_id: types.RequestId | None = None,
) -> None:
"""Send an elicitation completion notification.
This should be sent when a URL mode elicitation has been completed
out-of-band to inform the client that it may retry any requests
that were waiting for this elicitation.
Args:
elicitation_id: The unique identifier of the completed elicitation
related_request_id: Optional ID of the request that triggered this
"""
await self.send_notification(
types.ServerNotification(
types.ElicitCompleteNotification(
params=types.ElicitCompleteNotificationParams(elicitationId=elicitation_id)
)
),
related_request_id,
)
def _build_elicit_form_request(
self,
message: str,
requestedSchema: types.ElicitRequestedSchema,
related_task_id: str | None = None,
task: types.TaskMetadata | None = None,
) -> types.JSONRPCRequest:
"""Build a form mode elicitation request without sending it.
Args:
message: The message to present to the user
requestedSchema: Schema defining the expected response structure
related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata
task: If provided, makes this a task-augmented request
Returns:
A JSONRPCRequest ready to be sent or queued
"""
params = types.ElicitRequestFormParams(
message=message,
requestedSchema=requestedSchema,
task=task,
)
params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True)
# Add related-task metadata if associated with a parent task
if related_task_id is not None:
# Defensive: model_dump() never includes _meta, but guard against future changes
if "_meta" not in params_data: # pragma: no cover
params_data["_meta"] = {}
params_data["_meta"][RELATED_TASK_METADATA_KEY] = types.RelatedTaskMetadata(
taskId=related_task_id
).model_dump(by_alias=True)
request_id = f"task-{related_task_id}-{id(params)}" if related_task_id else self._request_id
if related_task_id is None:
self._request_id += 1
return types.JSONRPCRequest(
jsonrpc="2.0",
id=request_id,
method="elicitation/create",
params=params_data,
)
def _build_elicit_url_request(
self,
message: str,
url: str,
elicitation_id: str,
related_task_id: str | None = None,
) -> types.JSONRPCRequest:
"""Build a URL mode elicitation request without sending it.
Args:
message: Human-readable explanation of why the interaction is needed
url: The URL the user should navigate to
elicitation_id: Unique identifier for tracking this elicitation
related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata
Returns:
A JSONRPCRequest ready to be sent or queued
"""
params = types.ElicitRequestURLParams(
message=message,
url=url,
elicitationId=elicitation_id,
)
params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True)
# Add related-task metadata if associated with a parent task
if related_task_id is not None:
# Defensive: model_dump() never includes _meta, but guard against future changes
if "_meta" not in params_data: # pragma: no cover
params_data["_meta"] = {}
params_data["_meta"][RELATED_TASK_METADATA_KEY] = types.RelatedTaskMetadata(
taskId=related_task_id
).model_dump(by_alias=True)
request_id = f"task-{related_task_id}-{id(params)}" if related_task_id else self._request_id
if related_task_id is None:
self._request_id += 1
return types.JSONRPCRequest(
jsonrpc="2.0",
id=request_id,
method="elicitation/create",
params=params_data,
)
def _build_create_message_request(
self,
messages: list[types.SamplingMessage],
*,
max_tokens: int,
system_prompt: str | None = None,
include_context: types.IncludeContext | None = None,
temperature: float | None = None,
stop_sequences: list[str] | None = None,
metadata: dict[str, Any] | None = None,
model_preferences: types.ModelPreferences | None = None,
tools: list[types.Tool] | None = None,
tool_choice: types.ToolChoice | None = None,
related_task_id: str | None = None,
task: types.TaskMetadata | None = None,
) -> types.JSONRPCRequest:
"""Build a sampling/createMessage request without sending it.
Args:
messages: The conversation messages to send
max_tokens: Maximum number of tokens to generate
system_prompt: Optional system prompt
include_context: Optional context inclusion setting
temperature: Optional sampling temperature
stop_sequences: Optional stop sequences
metadata: Optional metadata to pass through to the LLM provider
model_preferences: Optional model selection preferences
tools: Optional list of tools the LLM can use during sampling
tool_choice: Optional control over tool usage behavior
related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata
task: If provided, makes this a task-augmented request
Returns:
A JSONRPCRequest ready to be sent or queued
"""
params = types.CreateMessageRequestParams(
messages=messages,
systemPrompt=system_prompt,
includeContext=include_context,
temperature=temperature,
maxTokens=max_tokens,
stopSequences=stop_sequences,
metadata=metadata,
modelPreferences=model_preferences,
tools=tools,
toolChoice=tool_choice,
task=task,
)
params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True)
# Add related-task metadata if associated with a parent task
if related_task_id is not None:
# Defensive: model_dump() never includes _meta, but guard against future changes
if "_meta" not in params_data: # pragma: no cover
params_data["_meta"] = {}
params_data["_meta"][RELATED_TASK_METADATA_KEY] = types.RelatedTaskMetadata(
taskId=related_task_id
).model_dump(by_alias=True)
request_id = f"task-{related_task_id}-{id(params)}" if related_task_id else self._request_id
if related_task_id is None:
self._request_id += 1
return types.JSONRPCRequest(
jsonrpc="2.0",
id=request_id,
method="sampling/createMessage",
params=params_data,
)
async def send_message(self, message: SessionMessage) -> None:
"""Send a raw session message.
This is primarily used by TaskResultHandler to deliver queued messages
(elicitation/sampling requests) to the client during task execution.
WARNING: This is a low-level experimental method that may change without
notice. Prefer using higher-level methods like send_notification() or
send_request() for normal operations.
Args:
message: The session message to send
"""
await self._write_stream.send(message)
async def _handle_incoming(self, req: ServerRequestResponder) -> None:
await self._incoming_message_stream_writer.send(req)
@property
def incoming_messages(
self,
) -> MemoryObjectReceiveStream[ServerRequestResponder]:
return self._incoming_message_stream_reader
@@ -0,0 +1,249 @@
"""
SSE Server Transport Module
This module implements a Server-Sent Events (SSE) transport layer for MCP servers.
Example usage:
```
# Create an SSE transport at an endpoint
sse = SseServerTransport("/messages/")
# Create Starlette routes for SSE and message handling
routes = [
Route("/sse", endpoint=handle_sse, methods=["GET"]),
Mount("/messages/", app=sse.handle_post_message),
]
# Define handler functions
async def handle_sse(request):
async with sse.connect_sse(
request.scope, request.receive, request._send
) as streams:
await app.run(
streams[0], streams[1], app.create_initialization_options()
)
# Return empty response to avoid NoneType error
return Response()
# Create and run Starlette app
starlette_app = Starlette(routes=routes)
uvicorn.run(starlette_app, host="127.0.0.1", port=port)
```
Note: The handle_sse function must return a Response to avoid a "TypeError: 'NoneType'
object is not callable" error when client disconnects. The example above returns
an empty Response() after the SSE connection ends to fix this.
See SseServerTransport class documentation for more details.
"""
import logging
from contextlib import asynccontextmanager
from typing import Any
from urllib.parse import quote
from uuid import UUID, uuid4
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import ValidationError
from sse_starlette import EventSourceResponse
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import Receive, Scope, Send
import mcp.types as types
from mcp.server.transport_security import (
TransportSecurityMiddleware,
TransportSecuritySettings,
)
from mcp.shared.message import ServerMessageMetadata, SessionMessage
logger = logging.getLogger(__name__)
class SseServerTransport:
"""
SSE server transport for MCP. This class provides _two_ ASGI applications,
suitable to be used with a framework like Starlette and a server like Hypercorn:
1. connect_sse() is an ASGI application which receives incoming GET requests,
and sets up a new SSE stream to send server messages to the client.
2. handle_post_message() is an ASGI application which receives incoming POST
requests, which should contain client messages that link to a
previously-established SSE session.
"""
_endpoint: str
_read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]]
_security: TransportSecurityMiddleware
def __init__(self, endpoint: str, security_settings: TransportSecuritySettings | None = None) -> None:
"""
Creates a new SSE server transport, which will direct the client to POST
messages to the relative path given.
Args:
endpoint: A relative path where messages should be posted
(e.g., "/messages/").
security_settings: Optional security settings for DNS rebinding protection.
Note:
We use relative paths instead of full URLs for several reasons:
1. Security: Prevents cross-origin requests by ensuring clients only connect
to the same origin they established the SSE connection with
2. Flexibility: The server can be mounted at any path without needing to
know its full URL
3. Portability: The same endpoint configuration works across different
environments (development, staging, production)
Raises:
ValueError: If the endpoint is a full URL instead of a relative path
"""
super().__init__()
# Validate that endpoint is a relative path and not a full URL
if "://" in endpoint or endpoint.startswith("//") or "?" in endpoint or "#" in endpoint:
raise ValueError(
f"Given endpoint: {endpoint} is not a relative path (e.g., '/messages/'), "
"expecting a relative path (e.g., '/messages/')."
)
# Ensure endpoint starts with a forward slash
if not endpoint.startswith("/"):
endpoint = "/" + endpoint
self._endpoint = endpoint
self._read_stream_writers = {}
self._security = TransportSecurityMiddleware(security_settings)
logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}")
@asynccontextmanager
async def connect_sse(self, scope: Scope, receive: Receive, send: Send): # pragma: no cover
if scope["type"] != "http":
logger.error("connect_sse received non-HTTP request")
raise ValueError("connect_sse can only handle HTTP requests")
# Validate request headers for DNS rebinding protection
request = Request(scope, receive)
error_response = await self._security.validate_request(request, is_post=False)
if error_response:
await error_response(scope, receive, send)
raise ValueError("Request validation failed")
logger.debug("Setting up SSE connection")
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
session_id = uuid4()
self._read_stream_writers[session_id] = read_stream_writer
logger.debug(f"Created new session with ID: {session_id}")
# Determine the full path for the message endpoint to be sent to the client.
# scope['root_path'] is the prefix where the current Starlette app
# instance is mounted.
# e.g., "" if top-level, or "/api_prefix" if mounted under "/api_prefix".
root_path = scope.get("root_path", "")
# self._endpoint is the path *within* this app, e.g., "/messages".
# Concatenating them gives the full absolute path from the server root.
# e.g., "" + "/messages" -> "/messages"
# e.g., "/api_prefix" + "/messages" -> "/api_prefix/messages"
full_message_path_for_client = root_path.rstrip("/") + self._endpoint
# This is the URI (path + query) the client will use to POST messages.
client_post_uri_data = f"{quote(full_message_path_for_client)}?session_id={session_id.hex}"
sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[dict[str, Any]](0)
async def sse_writer():
logger.debug("Starting SSE writer")
async with sse_stream_writer, write_stream_reader:
await sse_stream_writer.send({"event": "endpoint", "data": client_post_uri_data})
logger.debug(f"Sent endpoint event: {client_post_uri_data}")
async for session_message in write_stream_reader:
logger.debug(f"Sending message via SSE: {session_message}")
await sse_stream_writer.send(
{
"event": "message",
"data": session_message.message.model_dump_json(by_alias=True, exclude_none=True),
}
)
async with anyio.create_task_group() as tg:
async def response_wrapper(scope: Scope, receive: Receive, send: Send):
"""
The EventSourceResponse returning signals a client close / disconnect.
In this case we close our side of the streams to signal the client that
the connection has been closed.
"""
await EventSourceResponse(content=sse_stream_reader, data_sender_callable=sse_writer)(
scope, receive, send
)
await read_stream_writer.aclose()
await write_stream_reader.aclose()
logging.debug(f"Client session disconnected {session_id}")
logger.debug("Starting SSE response task")
tg.start_soon(response_wrapper, scope, receive, send)
logger.debug("Yielding read and write streams")
yield (read_stream, write_stream)
async def handle_post_message(self, scope: Scope, receive: Receive, send: Send) -> None: # pragma: no cover
logger.debug("Handling POST message")
request = Request(scope, receive)
# Validate request headers for DNS rebinding protection
error_response = await self._security.validate_request(request, is_post=True)
if error_response:
return await error_response(scope, receive, send)
session_id_param = request.query_params.get("session_id")
if session_id_param is None:
logger.warning("Received request without session_id")
response = Response("session_id is required", status_code=400)
return await response(scope, receive, send)
try:
session_id = UUID(hex=session_id_param)
logger.debug(f"Parsed session ID: {session_id}")
except ValueError:
logger.warning(f"Received invalid session ID: {session_id_param}")
response = Response("Invalid session ID", status_code=400)
return await response(scope, receive, send)
writer = self._read_stream_writers.get(session_id)
if not writer:
logger.warning(f"Could not find session for ID: {session_id}")
response = Response("Could not find session", status_code=404)
return await response(scope, receive, send)
body = await request.body()
logger.debug(f"Received JSON: {body}")
try:
message = types.JSONRPCMessage.model_validate_json(body)
logger.debug(f"Validated client message: {message}")
except ValidationError as err:
logger.exception("Failed to parse message")
response = Response("Could not parse message", status_code=400)
await response(scope, receive, send)
await writer.send(err)
return
# Pass the ASGI scope for framework-agnostic access to request data
metadata = ServerMessageMetadata(request_context=request)
session_message = SessionMessage(message, metadata=metadata)
logger.debug(f"Sending session message to writer: {session_message}")
response = Response("Accepted", status_code=202)
await response(scope, receive, send)
await writer.send(session_message)
@@ -0,0 +1,88 @@
"""
Stdio Server Transport Module
This module provides functionality for creating an stdio-based transport layer
that can be used to communicate with an MCP client through standard input/output
streams.
Example usage:
```
async def run_server():
async with stdio_server() as (read_stream, write_stream):
# read_stream contains incoming JSONRPCMessages from stdin
# write_stream allows sending JSONRPCMessages to stdout
server = await create_my_server()
await server.run(read_stream, write_stream, init_options)
anyio.run(run_server)
```
"""
import sys
from contextlib import asynccontextmanager
from io import TextIOWrapper
import anyio
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
import mcp.types as types
from mcp.shared.message import SessionMessage
@asynccontextmanager
async def stdio_server(
stdin: anyio.AsyncFile[str] | None = None,
stdout: anyio.AsyncFile[str] | None = None,
):
"""
Server transport for stdio: this communicates with an MCP client by reading
from the current process' stdin and writing to stdout.
"""
# Purposely not using context managers for these, as we don't want to close
# standard process handles. Encoding of stdin/stdout as text streams on
# python is platform-dependent (Windows is particularly problematic), so we
# re-wrap the underlying binary stream to ensure UTF-8.
if not stdin:
stdin = anyio.wrap_file(TextIOWrapper(sys.stdin.buffer, encoding="utf-8"))
if not stdout:
stdout = anyio.wrap_file(TextIOWrapper(sys.stdout.buffer, encoding="utf-8"))
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
async def stdin_reader():
try:
async with read_stream_writer:
async for line in stdin:
try:
message = types.JSONRPCMessage.model_validate_json(line)
except Exception as exc: # pragma: no cover
await read_stream_writer.send(exc)
continue
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
except anyio.ClosedResourceError: # pragma: no cover
await anyio.lowlevel.checkpoint()
async def stdout_writer():
try:
async with write_stream_reader:
async for session_message in write_stream_reader:
json = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
await stdout.write(json + "\n")
await stdout.flush()
except anyio.ClosedResourceError: # pragma: no cover
await anyio.lowlevel.checkpoint()
async with anyio.create_task_group() as tg:
tg.start_soon(stdin_reader)
tg.start_soon(stdout_writer)
yield read_stream, write_stream
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,297 @@
"""StreamableHTTP Session Manager for MCP servers."""
from __future__ import annotations
import contextlib
import logging
from collections.abc import AsyncIterator
from http import HTTPStatus
from typing import Any
from uuid import uuid4
import anyio
from anyio.abc import TaskStatus
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import Receive, Scope, Send
from mcp.server.lowlevel.server import Server as MCPServer
from mcp.server.streamable_http import (
MCP_SESSION_ID_HEADER,
EventStore,
StreamableHTTPServerTransport,
)
from mcp.server.transport_security import TransportSecuritySettings
from mcp.types import INVALID_REQUEST, ErrorData, JSONRPCError
logger = logging.getLogger(__name__)
class StreamableHTTPSessionManager:
"""
Manages StreamableHTTP sessions with optional resumability via event store.
This class abstracts away the complexity of session management, event storage,
and request handling for StreamableHTTP transports. It handles:
1. Session tracking for clients
2. Resumability via an optional event store
3. Connection management and lifecycle
4. Request handling and transport setup
Important: Only one StreamableHTTPSessionManager instance should be created
per application. The instance cannot be reused after its run() context has
completed. If you need to restart the manager, create a new instance.
Args:
app: The MCP server instance
event_store: Optional event store for resumability support.
If provided, enables resumable connections where clients
can reconnect and receive missed events.
If None, sessions are still tracked but not resumable.
json_response: Whether to use JSON responses instead of SSE streams
stateless: If True, creates a completely fresh transport for each request
with no session tracking or state persistence between requests.
security_settings: Optional transport security settings.
retry_interval: Retry interval in milliseconds to suggest to clients in SSE
retry field. Used for SSE polling behavior.
"""
def __init__(
self,
app: MCPServer[Any, Any],
event_store: EventStore | None = None,
json_response: bool = False,
stateless: bool = False,
security_settings: TransportSecuritySettings | None = None,
retry_interval: int | None = None,
):
self.app = app
self.event_store = event_store
self.json_response = json_response
self.stateless = stateless
self.security_settings = security_settings
self.retry_interval = retry_interval
# Session tracking (only used if not stateless)
self._session_creation_lock = anyio.Lock()
self._server_instances: dict[str, StreamableHTTPServerTransport] = {}
# The task group will be set during lifespan
self._task_group = None
# Thread-safe tracking of run() calls
self._run_lock = anyio.Lock()
self._has_started = False
@contextlib.asynccontextmanager
async def run(self) -> AsyncIterator[None]:
"""
Run the session manager with proper lifecycle management.
This creates and manages the task group for all session operations.
Important: This method can only be called once per instance. The same
StreamableHTTPSessionManager instance cannot be reused after this
context manager exits. Create a new instance if you need to restart.
Use this in the lifespan context manager of your Starlette app:
@contextlib.asynccontextmanager
async def lifespan(app: Starlette) -> AsyncIterator[None]:
async with session_manager.run():
yield
"""
# Thread-safe check to ensure run() is only called once
async with self._run_lock:
if self._has_started:
raise RuntimeError(
"StreamableHTTPSessionManager .run() can only be called "
"once per instance. Create a new instance if you need to run again."
)
self._has_started = True
async with anyio.create_task_group() as tg:
# Store the task group for later use
self._task_group = tg
logger.info("StreamableHTTP session manager started")
try:
yield # Let the application run
finally:
logger.info("StreamableHTTP session manager shutting down")
# Cancel task group to stop all spawned tasks
tg.cancel_scope.cancel()
self._task_group = None
# Clear any remaining server instances
self._server_instances.clear()
async def handle_request(
self,
scope: Scope,
receive: Receive,
send: Send,
) -> None:
"""
Process ASGI request with proper session handling and transport setup.
Dispatches to the appropriate handler based on stateless mode.
Args:
scope: ASGI scope
receive: ASGI receive function
send: ASGI send function
"""
if self._task_group is None:
raise RuntimeError("Task group is not initialized. Make sure to use run().")
# Dispatch to the appropriate handler
if self.stateless:
await self._handle_stateless_request(scope, receive, send)
else:
await self._handle_stateful_request(scope, receive, send)
async def _handle_stateless_request(
self,
scope: Scope,
receive: Receive,
send: Send,
) -> None:
"""
Process request in stateless mode - creating a new transport for each request.
Args:
scope: ASGI scope
receive: ASGI receive function
send: ASGI send function
"""
logger.debug("Stateless mode: Creating new transport for this request")
# No session ID needed in stateless mode
http_transport = StreamableHTTPServerTransport(
mcp_session_id=None, # No session tracking in stateless mode
is_json_response_enabled=self.json_response,
event_store=None, # No event store in stateless mode
security_settings=self.security_settings,
)
# Start server in a new task
async def run_stateless_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED):
async with http_transport.connect() as streams:
read_stream, write_stream = streams
task_status.started()
try:
await self.app.run(
read_stream,
write_stream,
self.app.create_initialization_options(),
stateless=True,
)
except Exception: # pragma: no cover
logger.exception("Stateless session crashed")
# Assert task group is not None for type checking
assert self._task_group is not None
# Start the server task
await self._task_group.start(run_stateless_server)
# Handle the HTTP request and return the response
await http_transport.handle_request(scope, receive, send)
# Terminate the transport after the request is handled
await http_transport.terminate()
async def _handle_stateful_request(
self,
scope: Scope,
receive: Receive,
send: Send,
) -> None:
"""
Process request in stateful mode - maintaining session state between requests.
Args:
scope: ASGI scope
receive: ASGI receive function
send: ASGI send function
"""
request = Request(scope, receive)
request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
# Existing session case
if request_mcp_session_id is not None and request_mcp_session_id in self._server_instances: # pragma: no cover
transport = self._server_instances[request_mcp_session_id]
logger.debug("Session already exists, handling request directly")
await transport.handle_request(scope, receive, send)
return
if request_mcp_session_id is None:
# New session case
logger.debug("Creating new transport")
async with self._session_creation_lock:
new_session_id = uuid4().hex
http_transport = StreamableHTTPServerTransport(
mcp_session_id=new_session_id,
is_json_response_enabled=self.json_response,
event_store=self.event_store, # May be None (no resumability)
security_settings=self.security_settings,
retry_interval=self.retry_interval,
)
assert http_transport.mcp_session_id is not None
self._server_instances[http_transport.mcp_session_id] = http_transport
logger.info(f"Created new transport with session ID: {new_session_id}")
# Define the server runner
async def run_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED) -> None:
async with http_transport.connect() as streams:
read_stream, write_stream = streams
task_status.started()
try:
await self.app.run(
read_stream,
write_stream,
self.app.create_initialization_options(),
stateless=False, # Stateful mode
)
except Exception as e:
logger.error(
f"Session {http_transport.mcp_session_id} crashed: {e}",
exc_info=True,
)
finally:
# Only remove from instances if not terminated
if ( # pragma: no branch
http_transport.mcp_session_id
and http_transport.mcp_session_id in self._server_instances
and not http_transport.is_terminated
):
logger.info(
"Cleaning up crashed session "
f"{http_transport.mcp_session_id} from "
"active instances."
)
del self._server_instances[http_transport.mcp_session_id]
# Assert task group is not None for type checking
assert self._task_group is not None
# Start the server task
await self._task_group.start(run_server)
# Handle the HTTP request and return the response
await http_transport.handle_request(scope, receive, send)
else:
# Unknown or expired session ID - return 404 per MCP spec
# TODO: Align error code once spec clarifies
# See: https://github.com/modelcontextprotocol/python-sdk/issues/1821
error_response = JSONRPCError(
jsonrpc="2.0",
id="server-error",
error=ErrorData(
code=INVALID_REQUEST,
message="Session not found",
),
)
response = Response(
content=error_response.model_dump_json(by_alias=True, exclude_none=True),
status_code=HTTPStatus.NOT_FOUND,
media_type="application/json",
)
await response(scope, receive, send)
@@ -0,0 +1,127 @@
"""DNS rebinding protection for MCP server transports."""
import logging
from pydantic import BaseModel, Field
from starlette.requests import Request
from starlette.responses import Response
logger = logging.getLogger(__name__)
class TransportSecuritySettings(BaseModel):
"""Settings for MCP transport security features.
These settings help protect against DNS rebinding attacks by validating
incoming request headers.
"""
enable_dns_rebinding_protection: bool = Field(
default=True,
description="Enable DNS rebinding protection (recommended for production)",
)
allowed_hosts: list[str] = Field(
default=[],
description="List of allowed Host header values. Only applies when "
+ "enable_dns_rebinding_protection is True.",
)
allowed_origins: list[str] = Field(
default=[],
description="List of allowed Origin header values. Only applies when "
+ "enable_dns_rebinding_protection is True.",
)
class TransportSecurityMiddleware:
"""Middleware to enforce DNS rebinding protection for MCP transport endpoints."""
def __init__(self, settings: TransportSecuritySettings | None = None):
# If not specified, disable DNS rebinding protection by default
# for backwards compatibility
self.settings = settings or TransportSecuritySettings(enable_dns_rebinding_protection=False)
def _validate_host(self, host: str | None) -> bool: # pragma: no cover
"""Validate the Host header against allowed values."""
if not host:
logger.warning("Missing Host header in request")
return False
# Check exact match first
if host in self.settings.allowed_hosts:
return True
# Check wildcard port patterns
for allowed in self.settings.allowed_hosts:
if allowed.endswith(":*"):
# Extract base host from pattern
base_host = allowed[:-2]
# Check if the actual host starts with base host and has a port
if host.startswith(base_host + ":"):
return True
logger.warning(f"Invalid Host header: {host}")
return False
def _validate_origin(self, origin: str | None) -> bool: # pragma: no cover
"""Validate the Origin header against allowed values."""
# Origin can be absent for same-origin requests
if not origin:
return True
# Check exact match first
if origin in self.settings.allowed_origins:
return True
# Check wildcard port patterns
for allowed in self.settings.allowed_origins:
if allowed.endswith(":*"):
# Extract base origin from pattern
base_origin = allowed[:-2]
# Check if the actual origin starts with base origin and has a port
if origin.startswith(base_origin + ":"):
return True
logger.warning(f"Invalid Origin header: {origin}")
return False
def _validate_content_type(self, content_type: str | None) -> bool: # pragma: no cover
"""Validate the Content-Type header for POST requests."""
if not content_type:
logger.warning("Missing Content-Type header in POST request")
return False
# Content-Type must start with application/json
if not content_type.lower().startswith("application/json"):
logger.warning(f"Invalid Content-Type header: {content_type}")
return False
return True
async def validate_request(self, request: Request, is_post: bool = False) -> Response | None:
"""Validate request headers for DNS rebinding protection.
Returns None if validation passes, or an error Response if validation fails.
"""
# Always validate Content-Type for POST requests
if is_post: # pragma: no branch
content_type = request.headers.get("content-type")
if not self._validate_content_type(content_type): # pragma: no cover
return Response("Invalid Content-Type header", status_code=400)
# Skip remaining validation if DNS rebinding protection is disabled
if not self.settings.enable_dns_rebinding_protection:
return None
# Validate Host header # pragma: no cover
host = request.headers.get("host") # pragma: no cover
if not self._validate_host(host): # pragma: no cover
return Response("Invalid Host header", status_code=421) # pragma: no cover
# Validate Origin header # pragma: no cover
origin = request.headers.get("origin") # pragma: no cover
if not self._validate_origin(origin): # pragma: no cover
return Response("Invalid Origin header", status_code=403) # pragma: no cover
return None # pragma: no cover
@@ -0,0 +1,104 @@
"""
Shared validation functions for server requests.
This module provides validation logic for sampling and elicitation requests
that is shared across normal and task-augmented code paths.
"""
from mcp.shared.exceptions import McpError
from mcp.types import (
INVALID_PARAMS,
ClientCapabilities,
ErrorData,
SamplingMessage,
Tool,
ToolChoice,
)
def check_sampling_tools_capability(client_caps: ClientCapabilities | None) -> bool:
"""
Check if the client supports sampling tools capability.
Args:
client_caps: The client's declared capabilities
Returns:
True if client supports sampling.tools, False otherwise
"""
if client_caps is None:
return False
if client_caps.sampling is None:
return False
if client_caps.sampling.tools is None:
return False
return True
def validate_sampling_tools(
client_caps: ClientCapabilities | None,
tools: list[Tool] | None,
tool_choice: ToolChoice | None,
) -> None:
"""
Validate that the client supports sampling tools if tools are being used.
Args:
client_caps: The client's declared capabilities
tools: The tools list, if provided
tool_choice: The tool choice setting, if provided
Raises:
McpError: If tools/tool_choice are provided but client doesn't support them
"""
if tools is not None or tool_choice is not None:
if not check_sampling_tools_capability(client_caps):
raise McpError(
ErrorData(
code=INVALID_PARAMS,
message="Client does not support sampling tools capability",
)
)
def validate_tool_use_result_messages(messages: list[SamplingMessage]) -> None:
"""
Validate tool_use/tool_result message structure per SEP-1577.
This validation ensures:
1. Messages with tool_result content contain ONLY tool_result content
2. tool_result messages are preceded by a message with tool_use
3. tool_result IDs match the tool_use IDs from the previous message
See: https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1577
Args:
messages: The list of sampling messages to validate
Raises:
ValueError: If the message structure is invalid
"""
if not messages:
return
last_content = messages[-1].content_as_list
has_tool_results = any(c.type == "tool_result" for c in last_content)
previous_content = messages[-2].content_as_list if len(messages) >= 2 else None
has_previous_tool_use = previous_content and any(c.type == "tool_use" for c in previous_content)
if has_tool_results:
# Per spec: "SamplingMessage with tool result content blocks
# MUST NOT contain other content types."
if any(c.type != "tool_result" for c in last_content):
raise ValueError("The last message must contain only tool_result content if any is present")
if previous_content is None:
raise ValueError("tool_result requires a previous message containing tool_use")
if not has_previous_tool_use:
raise ValueError("tool_result blocks do not match any tool_use in the previous message")
if has_previous_tool_use and previous_content:
tool_use_ids = {c.id for c in previous_content if c.type == "tool_use"}
tool_result_ids = {c.toolUseId for c in last_content if c.type == "tool_result"}
if tool_use_ids != tool_result_ids:
raise ValueError("ids of tool_result blocks and tool_use blocks from previous message do not match")
@@ -0,0 +1,62 @@
import logging
from contextlib import asynccontextmanager
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic_core import ValidationError
from starlette.types import Receive, Scope, Send
from starlette.websockets import WebSocket
import mcp.types as types
from mcp.shared.message import SessionMessage
logger = logging.getLogger(__name__)
@asynccontextmanager # pragma: no cover
async def websocket_server(scope: Scope, receive: Receive, send: Send):
"""
WebSocket server transport for MCP. This is an ASGI application, suitable to be
used with a framework like Starlette and a server like Hypercorn.
"""
websocket = WebSocket(scope, receive, send)
await websocket.accept(subprotocol="mcp")
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
async def ws_reader():
try:
async with read_stream_writer:
async for msg in websocket.iter_text():
try:
client_message = types.JSONRPCMessage.model_validate_json(msg)
except ValidationError as exc:
await read_stream_writer.send(exc)
continue
session_message = SessionMessage(client_message)
await read_stream_writer.send(session_message)
except anyio.ClosedResourceError:
await websocket.close()
async def ws_writer():
try:
async with write_stream_reader:
async for session_message in write_stream_reader:
obj = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
await websocket.send_text(obj)
except anyio.ClosedResourceError:
await websocket.close()
async with anyio.create_task_group() as tg:
tg.start_soon(ws_reader)
tg.start_soon(ws_writer)
yield (read_stream, write_stream)
Whitespace-only changes.
@@ -0,0 +1,87 @@
"""Utilities for creating standardized httpx AsyncClient instances."""
from typing import Any, Protocol
import httpx
__all__ = ["create_mcp_http_client", "MCP_DEFAULT_TIMEOUT", "MCP_DEFAULT_SSE_READ_TIMEOUT"]
# Default MCP timeout configuration
MCP_DEFAULT_TIMEOUT = 30.0 # General operations (seconds)
MCP_DEFAULT_SSE_READ_TIMEOUT = 300.0 # SSE streams - 5 minutes (seconds)
class McpHttpClientFactory(Protocol): # pragma: no branch
def __call__( # pragma: no branch
self,
headers: dict[str, str] | None = None,
timeout: httpx.Timeout | None = None,
auth: httpx.Auth | None = None,
) -> httpx.AsyncClient: ...
def create_mcp_http_client(
headers: dict[str, str] | None = None,
timeout: httpx.Timeout | None = None,
auth: httpx.Auth | None = None,
) -> httpx.AsyncClient:
"""Create a standardized httpx AsyncClient with MCP defaults.
This function provides common defaults used throughout the MCP codebase:
- follow_redirects=True (always enabled)
- Default timeout of 30 seconds if not specified
Args:
headers: Optional headers to include with all requests.
timeout: Request timeout as httpx.Timeout object.
Defaults to 30 seconds if not specified.
auth: Optional authentication handler.
Returns:
Configured httpx.AsyncClient instance with MCP defaults.
Note:
The returned AsyncClient must be used as a context manager to ensure
proper cleanup of connections.
Examples:
# Basic usage with MCP defaults
async with create_mcp_http_client() as client:
response = await client.get("https://api.example.com")
# With custom headers
headers = {"Authorization": "Bearer token"}
async with create_mcp_http_client(headers) as client:
response = await client.get("/endpoint")
# With both custom headers and timeout
timeout = httpx.Timeout(60.0, read=300.0)
async with create_mcp_http_client(headers, timeout) as client:
response = await client.get("/long-request")
# With authentication
from httpx import BasicAuth
auth = BasicAuth(username="user", password="pass")
async with create_mcp_http_client(headers, timeout, auth) as client:
response = await client.get("/protected-endpoint")
"""
# Set MCP defaults
kwargs: dict[str, Any] = {
"follow_redirects": True,
}
# Handle timeout
if timeout is None:
kwargs["timeout"] = httpx.Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT)
else:
kwargs["timeout"] = timeout
# Handle headers
if headers is not None:
kwargs["headers"] = headers
# Handle authentication
if auth is not None: # pragma: no cover
kwargs["auth"] = auth
return httpx.AsyncClient(**kwargs)
@@ -0,0 +1,159 @@
from typing import Any, Literal
from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, field_validator
class OAuthToken(BaseModel):
"""
See https://datatracker.ietf.org/doc/html/rfc6749#section-5.1
"""
access_token: str
token_type: Literal["Bearer"] = "Bearer"
expires_in: int | None = None
scope: str | None = None
refresh_token: str | None = None
@field_validator("token_type", mode="before")
@classmethod
def normalize_token_type(cls, v: str | None) -> str | None:
if isinstance(v, str):
# Bearer is title-cased in the spec, so we normalize it
# https://datatracker.ietf.org/doc/html/rfc6750#section-4
return v.title()
return v # pragma: no cover
class InvalidScopeError(Exception):
def __init__(self, message: str):
self.message = message
class InvalidRedirectUriError(Exception):
def __init__(self, message: str):
self.message = message
class OAuthClientMetadata(BaseModel):
"""
RFC 7591 OAuth 2.0 Dynamic Client Registration metadata.
See https://datatracker.ietf.org/doc/html/rfc7591#section-2
for the full specification.
"""
redirect_uris: list[AnyUrl] | None = Field(..., min_length=1)
# supported auth methods for the token endpoint
token_endpoint_auth_method: (
Literal["none", "client_secret_post", "client_secret_basic", "private_key_jwt"] | None
) = None
# supported grant_types of this implementation
grant_types: list[
Literal["authorization_code", "refresh_token", "urn:ietf:params:oauth:grant-type:jwt-bearer"] | str
] = [
"authorization_code",
"refresh_token",
]
# The MCP spec requires the "code" response type, but OAuth
# servers may also return additional types they support
response_types: list[str] = ["code"]
scope: str | None = None
# these fields are currently unused, but we support & store them for potential
# future use
client_name: str | None = None
client_uri: AnyHttpUrl | None = None
logo_uri: AnyHttpUrl | None = None
contacts: list[str] | None = None
tos_uri: AnyHttpUrl | None = None
policy_uri: AnyHttpUrl | None = None
jwks_uri: AnyHttpUrl | None = None
jwks: Any | None = None
software_id: str | None = None
software_version: str | None = None
def validate_scope(self, requested_scope: str | None) -> list[str] | None:
if requested_scope is None:
return None
requested_scopes = requested_scope.split(" ")
allowed_scopes = [] if self.scope is None else self.scope.split(" ")
for scope in requested_scopes:
if scope not in allowed_scopes: # pragma: no branch
raise InvalidScopeError(f"Client was not registered with scope {scope}")
return requested_scopes # pragma: no cover
def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl:
if redirect_uri is not None:
# Validate redirect_uri against client's registered redirect URIs
if self.redirect_uris is None or redirect_uri not in self.redirect_uris:
raise InvalidRedirectUriError(f"Redirect URI '{redirect_uri}' not registered for client")
return redirect_uri
elif self.redirect_uris is not None and len(self.redirect_uris) == 1:
return self.redirect_uris[0]
else:
raise InvalidRedirectUriError("redirect_uri must be specified when client has multiple registered URIs")
class OAuthClientInformationFull(OAuthClientMetadata):
"""
RFC 7591 OAuth 2.0 Dynamic Client Registration full response
(client information plus metadata).
"""
client_id: str | None = None
client_secret: str | None = None
client_id_issued_at: int | None = None
client_secret_expires_at: int | None = None
class OAuthMetadata(BaseModel):
"""
RFC 8414 OAuth 2.0 Authorization Server Metadata.
See https://datatracker.ietf.org/doc/html/rfc8414#section-2
"""
issuer: AnyHttpUrl
authorization_endpoint: AnyHttpUrl
token_endpoint: AnyHttpUrl
registration_endpoint: AnyHttpUrl | None = None
scopes_supported: list[str] | None = None
response_types_supported: list[str] = ["code"]
response_modes_supported: list[str] | None = None
grant_types_supported: list[str] | None = None
token_endpoint_auth_methods_supported: list[str] | None = None
token_endpoint_auth_signing_alg_values_supported: list[str] | None = None
service_documentation: AnyHttpUrl | None = None
ui_locales_supported: list[str] | None = None
op_policy_uri: AnyHttpUrl | None = None
op_tos_uri: AnyHttpUrl | None = None
revocation_endpoint: AnyHttpUrl | None = None
revocation_endpoint_auth_methods_supported: list[str] | None = None
revocation_endpoint_auth_signing_alg_values_supported: list[str] | None = None
introspection_endpoint: AnyHttpUrl | None = None
introspection_endpoint_auth_methods_supported: list[str] | None = None
introspection_endpoint_auth_signing_alg_values_supported: list[str] | None = None
code_challenge_methods_supported: list[str] | None = None
client_id_metadata_document_supported: bool | None = None
class ProtectedResourceMetadata(BaseModel):
"""
RFC 9728 OAuth 2.0 Protected Resource Metadata.
See https://datatracker.ietf.org/doc/html/rfc9728#section-2
"""
resource: AnyHttpUrl
authorization_servers: list[AnyHttpUrl] = Field(..., min_length=1)
jwks_uri: AnyHttpUrl | None = None
scopes_supported: list[str] | None = None
bearer_methods_supported: list[str] | None = Field(default=["header"]) # MCP only supports header method
resource_signing_alg_values_supported: list[str] | None = None
resource_name: str | None = None
resource_documentation: AnyHttpUrl | None = None
resource_policy_uri: AnyHttpUrl | None = None
resource_tos_uri: AnyHttpUrl | None = None
# tls_client_certificate_bound_access_tokens default is False, but ommited here for clarity
tls_client_certificate_bound_access_tokens: bool | None = None
authorization_details_types_supported: list[str] | None = None
dpop_signing_alg_values_supported: list[str] | None = None
# dpop_bound_access_tokens_required default is False, but ommited here for clarity
dpop_bound_access_tokens_required: bool | None = None
@@ -0,0 +1,85 @@
"""Utilities for OAuth 2.0 Resource Indicators (RFC 8707) and PKCE (RFC 7636)."""
import time
from urllib.parse import urlparse, urlsplit, urlunsplit
from pydantic import AnyUrl, HttpUrl
def resource_url_from_server_url(url: str | HttpUrl | AnyUrl) -> str:
"""Convert server URL to canonical resource URL per RFC 8707.
RFC 8707 section 2 states that resource URIs "MUST NOT include a fragment component".
Returns absolute URI with lowercase scheme/host for canonical form.
Args:
url: Server URL to convert
Returns:
Canonical resource URL string
"""
# Convert to string if needed
url_str = str(url)
# Parse the URL and remove fragment, create canonical form
parsed = urlsplit(url_str)
canonical = urlunsplit(parsed._replace(scheme=parsed.scheme.lower(), netloc=parsed.netloc.lower(), fragment=""))
return canonical
def check_resource_allowed(requested_resource: str, configured_resource: str) -> bool:
"""Check if a requested resource URL matches a configured resource URL.
A requested resource matches if it has the same scheme, domain, port,
and its path starts with the configured resource's path. This allows
hierarchical matching where a token for a parent resource can be used
for child resources.
Args:
requested_resource: The resource URL being requested
configured_resource: The resource URL that has been configured
Returns:
True if the requested resource matches the configured resource
"""
# Parse both URLs
requested = urlparse(requested_resource)
configured = urlparse(configured_resource)
# Compare scheme, host, and port (origin)
if requested.scheme.lower() != configured.scheme.lower() or requested.netloc.lower() != configured.netloc.lower():
return False
# Handle cases like requested=/foo and configured=/foo/
requested_path = requested.path
configured_path = configured.path
# If requested path is shorter, it cannot be a child
if len(requested_path) < len(configured_path):
return False
# Check if the requested path starts with the configured path
# Ensure both paths end with / for proper comparison
# This ensures that paths like "/api123" don't incorrectly match "/api"
if not requested_path.endswith("/"):
requested_path += "/"
if not configured_path.endswith("/"):
configured_path += "/"
return requested_path.startswith(configured_path)
def calculate_token_expiry(expires_in: int | str | None) -> float | None:
"""Calculate token expiry timestamp from expires_in seconds.
Args:
expires_in: Seconds until token expiration (may be string from some servers)
Returns:
Unix timestamp when token expires, or None if no expiry specified
"""
if expires_in is None:
return None # pragma: no cover
# Defensive: handle servers that return expires_in as string
return time.time() + int(expires_in)
@@ -0,0 +1,32 @@
"""
Request context for MCP handlers.
"""
from dataclasses import dataclass, field
from typing import Any, Generic
from typing_extensions import TypeVar
from mcp.shared.message import CloseSSEStreamCallback
from mcp.shared.session import BaseSession
from mcp.types import RequestId, RequestParams
SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
LifespanContextT = TypeVar("LifespanContextT")
RequestT = TypeVar("RequestT", default=Any)
@dataclass
class RequestContext(Generic[SessionT, LifespanContextT, RequestT]):
request_id: RequestId
meta: RequestParams.Meta | None
session: SessionT
lifespan_context: LifespanContextT
# NOTE: This is typed as Any to avoid circular imports. The actual type is
# mcp.server.experimental.request_context.Experimental, but importing it here
# triggers mcp.server.__init__ -> fastmcp -> tools -> back to this module.
# The Server sets this to an Experimental instance at runtime.
experimental: Any = field(default=None)
request: RequestT | None = None
close_sse_stream: CloseSSEStreamCallback | None = None
close_standalone_sse_stream: CloseSSEStreamCallback | None = None
@@ -0,0 +1,71 @@
from __future__ import annotations
from typing import Any, cast
from mcp.types import URL_ELICITATION_REQUIRED, ElicitRequestURLParams, ErrorData
class McpError(Exception):
"""
Exception type raised when an error arrives over an MCP connection.
"""
error: ErrorData
def __init__(self, error: ErrorData):
"""Initialize McpError."""
super().__init__(error.message)
self.error = error
class UrlElicitationRequiredError(McpError):
"""
Specialized error for when a tool requires URL mode elicitation(s) before proceeding.
Servers can raise this error from tool handlers to indicate that the client
must complete one or more URL elicitations before the request can be processed.
Example:
raise UrlElicitationRequiredError([
ElicitRequestURLParams(
mode="url",
message="Authorization required for your files",
url="https://example.com/oauth/authorize",
elicitationId="auth-001"
)
])
"""
def __init__(
self,
elicitations: list[ElicitRequestURLParams],
message: str | None = None,
):
"""Initialize UrlElicitationRequiredError."""
if message is None:
message = f"URL elicitation{'s' if len(elicitations) > 1 else ''} required"
self._elicitations = elicitations
error = ErrorData(
code=URL_ELICITATION_REQUIRED,
message=message,
data={"elicitations": [e.model_dump(by_alias=True, exclude_none=True) for e in elicitations]},
)
super().__init__(error)
@property
def elicitations(self) -> list[ElicitRequestURLParams]:
"""The list of URL elicitations required before the request can proceed."""
return self._elicitations
@classmethod
def from_error(cls, error: ErrorData) -> UrlElicitationRequiredError:
"""Reconstruct from an ErrorData received over the wire."""
if error.code != URL_ELICITATION_REQUIRED:
raise ValueError(f"Expected error code {URL_ELICITATION_REQUIRED}, got {error.code}")
data = cast(dict[str, Any], error.data or {})
raw_elicitations = cast(list[dict[str, Any]], data.get("elicitations", []))
elicitations = [ElicitRequestURLParams.model_validate(e) for e in raw_elicitations]
return cls(elicitations, error.message)
@@ -0,0 +1,7 @@
"""
Pure experimental MCP features (no server dependencies).
WARNING: These APIs are experimental and may change without notice.
For server-integrated experimental features, use mcp.server.experimental.
"""
@@ -0,0 +1,12 @@
"""
Pure task state management for MCP.
WARNING: These APIs are experimental and may change without notice.
Import directly from submodules:
- mcp.shared.experimental.tasks.store.TaskStore
- mcp.shared.experimental.tasks.context.TaskContext
- mcp.shared.experimental.tasks.in_memory_task_store.InMemoryTaskStore
- mcp.shared.experimental.tasks.message_queue.TaskMessageQueue
- mcp.shared.experimental.tasks.helpers.is_terminal
"""
@@ -0,0 +1,115 @@
"""
Tasks capability checking utilities.
This module provides functions for checking and requiring task-related
capabilities. All tasks capability logic is centralized here to keep
the main session code clean.
WARNING: These APIs are experimental and may change without notice.
"""
from mcp.shared.exceptions import McpError
from mcp.types import (
INVALID_REQUEST,
ClientCapabilities,
ClientTasksCapability,
ErrorData,
)
def check_tasks_capability(
required: ClientTasksCapability,
client: ClientTasksCapability,
) -> bool:
"""
Check if client's tasks capability matches the required capability.
Args:
required: The capability being checked for
client: The client's declared capabilities
Returns:
True if client has the required capability, False otherwise
"""
if required.requests is None:
return True
if client.requests is None:
return False
# Check elicitation.create
if required.requests.elicitation is not None:
if client.requests.elicitation is None:
return False
if required.requests.elicitation.create is not None:
if client.requests.elicitation.create is None:
return False
# Check sampling.createMessage
if required.requests.sampling is not None:
if client.requests.sampling is None:
return False
if required.requests.sampling.createMessage is not None:
if client.requests.sampling.createMessage is None:
return False
return True
def has_task_augmented_elicitation(caps: ClientCapabilities) -> bool:
"""Check if capabilities include task-augmented elicitation support."""
if caps.tasks is None:
return False
if caps.tasks.requests is None:
return False
if caps.tasks.requests.elicitation is None:
return False
return caps.tasks.requests.elicitation.create is not None
def has_task_augmented_sampling(caps: ClientCapabilities) -> bool:
"""Check if capabilities include task-augmented sampling support."""
if caps.tasks is None:
return False
if caps.tasks.requests is None:
return False
if caps.tasks.requests.sampling is None:
return False
return caps.tasks.requests.sampling.createMessage is not None
def require_task_augmented_elicitation(client_caps: ClientCapabilities | None) -> None:
"""
Raise McpError if client doesn't support task-augmented elicitation.
Args:
client_caps: The client's declared capabilities, or None if not initialized
Raises:
McpError: If client doesn't support task-augmented elicitation
"""
if client_caps is None or not has_task_augmented_elicitation(client_caps):
raise McpError(
ErrorData(
code=INVALID_REQUEST,
message="Client does not support task-augmented elicitation",
)
)
def require_task_augmented_sampling(client_caps: ClientCapabilities | None) -> None:
"""
Raise McpError if client doesn't support task-augmented sampling.
Args:
client_caps: The client's declared capabilities, or None if not initialized
Raises:
McpError: If client doesn't support task-augmented sampling
"""
if client_caps is None or not has_task_augmented_sampling(client_caps):
raise McpError(
ErrorData(
code=INVALID_REQUEST,
message="Client does not support task-augmented sampling",
)
)
@@ -0,0 +1,101 @@
"""
TaskContext - Pure task state management.
This module provides TaskContext, which manages task state without any
server/session dependencies. It can be used standalone for distributed
workers or wrapped by ServerTaskContext for full server integration.
"""
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.types import TASK_STATUS_COMPLETED, TASK_STATUS_FAILED, Result, Task
class TaskContext:
"""
Pure task state management - no session dependencies.
This class handles:
- Task state (status, result)
- Cancellation tracking
- Store interactions
For server-integrated features (elicit, create_message, notifications),
use ServerTaskContext from mcp.server.experimental.
Example (distributed worker):
async def worker_job(task_id: str):
store = RedisTaskStore(redis_url)
task = await store.get_task(task_id)
ctx = TaskContext(task=task, store=store)
await ctx.update_status("Working...")
result = await do_work()
await ctx.complete(result)
"""
def __init__(self, task: Task, store: TaskStore):
self._task = task
self._store = store
self._cancelled = False
@property
def task_id(self) -> str:
"""The task identifier."""
return self._task.taskId
@property
def task(self) -> Task:
"""The current task state."""
return self._task
@property
def is_cancelled(self) -> bool:
"""Whether cancellation has been requested."""
return self._cancelled
def request_cancellation(self) -> None:
"""
Request cancellation of this task.
This sets is_cancelled=True. Task work should check this
periodically and exit gracefully if set.
"""
self._cancelled = True
async def update_status(self, message: str) -> None:
"""
Update the task's status message.
Args:
message: The new status message
"""
self._task = await self._store.update_task(
self.task_id,
status_message=message,
)
async def complete(self, result: Result) -> None:
"""
Mark the task as completed with the given result.
Args:
result: The task result
"""
await self._store.store_result(self.task_id, result)
self._task = await self._store.update_task(
self.task_id,
status=TASK_STATUS_COMPLETED,
)
async def fail(self, error: str) -> None:
"""
Mark the task as failed with an error message.
Args:
error: The error message
"""
self._task = await self._store.update_task(
self.task_id,
status=TASK_STATUS_FAILED,
status_message=error,
)
@@ -0,0 +1,181 @@
"""
Helper functions for pure task management.
These helpers work with pure TaskContext and don't require server dependencies.
For server-integrated task helpers, use mcp.server.experimental.
"""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from uuid import uuid4
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.context import TaskContext
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.types import (
INVALID_PARAMS,
TASK_STATUS_CANCELLED,
TASK_STATUS_COMPLETED,
TASK_STATUS_FAILED,
TASK_STATUS_WORKING,
CancelTaskResult,
ErrorData,
Task,
TaskMetadata,
TaskStatus,
)
# Metadata key for model-immediate-response (per MCP spec)
# Servers MAY include this in CreateTaskResult._meta to provide an immediate
# response string while the task executes in the background.
MODEL_IMMEDIATE_RESPONSE_KEY = "io.modelcontextprotocol/model-immediate-response"
# Metadata key for associating requests with a task (per MCP spec)
RELATED_TASK_METADATA_KEY = "io.modelcontextprotocol/related-task"
def is_terminal(status: TaskStatus) -> bool:
"""
Check if a task status represents a terminal state.
Terminal states are those where the task has finished and will not change.
Args:
status: The task status to check
Returns:
True if the status is terminal (completed, failed, or cancelled)
"""
return status in (TASK_STATUS_COMPLETED, TASK_STATUS_FAILED, TASK_STATUS_CANCELLED)
async def cancel_task(
store: TaskStore,
task_id: str,
) -> CancelTaskResult:
"""
Cancel a task with spec-compliant validation.
Per spec: "Receivers MUST reject cancellation of terminal status tasks
with -32602 (Invalid params)"
This helper validates that the task exists and is not in a terminal state
before setting it to "cancelled".
Args:
store: The task store
task_id: The task identifier to cancel
Returns:
CancelTaskResult with the cancelled task state
Raises:
McpError: With INVALID_PARAMS (-32602) if:
- Task does not exist
- Task is already in a terminal state (completed, failed, cancelled)
Example:
@server.experimental.cancel_task()
async def handle_cancel(request: CancelTaskRequest) -> CancelTaskResult:
return await cancel_task(store, request.params.taskId)
"""
task = await store.get_task(task_id)
if task is None:
raise McpError(
ErrorData(
code=INVALID_PARAMS,
message=f"Task not found: {task_id}",
)
)
if is_terminal(task.status):
raise McpError(
ErrorData(
code=INVALID_PARAMS,
message=f"Cannot cancel task in terminal state '{task.status}'",
)
)
# Update task to cancelled status
cancelled_task = await store.update_task(task_id, status=TASK_STATUS_CANCELLED)
return CancelTaskResult(**cancelled_task.model_dump())
def generate_task_id() -> str:
"""Generate a unique task ID."""
return str(uuid4())
def create_task_state(
metadata: TaskMetadata,
task_id: str | None = None,
) -> Task:
"""
Create a Task object with initial state.
This is a helper for TaskStore implementations.
Args:
metadata: Task metadata
task_id: Optional task ID (generated if not provided)
Returns:
A new Task in "working" status
"""
now = datetime.now(timezone.utc)
return Task(
taskId=task_id or generate_task_id(),
status=TASK_STATUS_WORKING,
createdAt=now,
lastUpdatedAt=now,
ttl=metadata.ttl,
pollInterval=500, # Default 500ms poll interval
)
@asynccontextmanager
async def task_execution(
task_id: str,
store: TaskStore,
) -> AsyncIterator[TaskContext]:
"""
Context manager for safe task execution (pure, no server dependencies).
Loads a task from the store and provides a TaskContext for the work.
If an unhandled exception occurs, the task is automatically marked as failed
and the exception is suppressed (since the failure is captured in task state).
This is useful for distributed workers that don't have a server session.
Args:
task_id: The task identifier to execute
store: The task store (must be accessible by the worker)
Yields:
TaskContext for updating status and completing/failing the task
Raises:
ValueError: If the task is not found in the store
Example (distributed worker):
async def worker_process(task_id: str):
store = RedisTaskStore(redis_url)
async with task_execution(task_id, store) as ctx:
await ctx.update_status("Working...")
result = await do_work()
await ctx.complete(result)
"""
task = await store.get_task(task_id)
if task is None:
raise ValueError(f"Task {task_id} not found")
ctx = TaskContext(task, store)
try:
yield ctx
except Exception as e:
# Auto-fail the task if an exception occurs and task isn't already terminal
# Exception is suppressed since failure is captured in task state
if not is_terminal(ctx.task.status):
await ctx.fail(str(e))
# Don't re-raise - the failure is recorded in task state
@@ -0,0 +1,219 @@
"""
In-memory implementation of TaskStore for demonstration purposes.
This implementation stores all tasks in memory and provides automatic cleanup
based on the TTL duration specified in the task metadata using lazy expiration.
Note: This is not suitable for production use as all data is lost on restart.
For production, consider implementing TaskStore with a database or distributed cache.
"""
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
import anyio
from mcp.shared.experimental.tasks.helpers import create_task_state, is_terminal
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.types import Result, Task, TaskMetadata, TaskStatus
@dataclass
class StoredTask:
"""Internal storage representation of a task."""
task: Task
result: Result | None = None
# Time when this task should be removed (None = never)
expires_at: datetime | None = field(default=None)
class InMemoryTaskStore(TaskStore):
"""
A simple in-memory implementation of TaskStore.
Features:
- Automatic TTL-based cleanup (lazy expiration)
- Thread-safe for single-process async use
- Pagination support for list_tasks
Limitations:
- All data lost on restart
- Not suitable for distributed systems
- No persistence
For production, implement TaskStore with Redis, PostgreSQL, etc.
"""
def __init__(self, page_size: int = 10) -> None:
self._tasks: dict[str, StoredTask] = {}
self._page_size = page_size
self._update_events: dict[str, anyio.Event] = {}
def _calculate_expiry(self, ttl_ms: int | None) -> datetime | None:
"""Calculate expiry time from TTL in milliseconds."""
if ttl_ms is None:
return None
return datetime.now(timezone.utc) + timedelta(milliseconds=ttl_ms)
def _is_expired(self, stored: StoredTask) -> bool:
"""Check if a task has expired."""
if stored.expires_at is None:
return False
return datetime.now(timezone.utc) >= stored.expires_at
def _cleanup_expired(self) -> None:
"""Remove all expired tasks. Called lazily during access operations."""
expired_ids = [task_id for task_id, stored in self._tasks.items() if self._is_expired(stored)]
for task_id in expired_ids:
del self._tasks[task_id]
async def create_task(
self,
metadata: TaskMetadata,
task_id: str | None = None,
) -> Task:
"""Create a new task with the given metadata."""
# Cleanup expired tasks on access
self._cleanup_expired()
task = create_task_state(metadata, task_id)
if task.taskId in self._tasks:
raise ValueError(f"Task with ID {task.taskId} already exists")
stored = StoredTask(
task=task,
expires_at=self._calculate_expiry(metadata.ttl),
)
self._tasks[task.taskId] = stored
# Return a copy to prevent external modification
return Task(**task.model_dump())
async def get_task(self, task_id: str) -> Task | None:
"""Get a task by ID."""
# Cleanup expired tasks on access
self._cleanup_expired()
stored = self._tasks.get(task_id)
if stored is None:
return None
# Return a copy to prevent external modification
return Task(**stored.task.model_dump())
async def update_task(
self,
task_id: str,
status: TaskStatus | None = None,
status_message: str | None = None,
) -> Task:
"""Update a task's status and/or message."""
stored = self._tasks.get(task_id)
if stored is None:
raise ValueError(f"Task with ID {task_id} not found")
# Per spec: Terminal states MUST NOT transition to any other status
if status is not None and status != stored.task.status and is_terminal(stored.task.status):
raise ValueError(f"Cannot transition from terminal status '{stored.task.status}'")
status_changed = False
if status is not None and stored.task.status != status:
stored.task.status = status
status_changed = True
if status_message is not None:
stored.task.statusMessage = status_message
# Update lastUpdatedAt on any change
stored.task.lastUpdatedAt = datetime.now(timezone.utc)
# If task is now terminal and has TTL, reset expiry timer
if status is not None and is_terminal(status) and stored.task.ttl is not None:
stored.expires_at = self._calculate_expiry(stored.task.ttl)
# Notify waiters if status changed
if status_changed:
await self.notify_update(task_id)
return Task(**stored.task.model_dump())
async def store_result(self, task_id: str, result: Result) -> None:
"""Store the result for a task."""
stored = self._tasks.get(task_id)
if stored is None:
raise ValueError(f"Task with ID {task_id} not found")
stored.result = result
async def get_result(self, task_id: str) -> Result | None:
"""Get the stored result for a task."""
stored = self._tasks.get(task_id)
if stored is None:
return None
return stored.result
async def list_tasks(
self,
cursor: str | None = None,
) -> tuple[list[Task], str | None]:
"""List tasks with pagination."""
# Cleanup expired tasks on access
self._cleanup_expired()
all_task_ids = list(self._tasks.keys())
start_index = 0
if cursor is not None:
try:
cursor_index = all_task_ids.index(cursor)
start_index = cursor_index + 1
except ValueError:
raise ValueError(f"Invalid cursor: {cursor}")
page_task_ids = all_task_ids[start_index : start_index + self._page_size]
tasks = [Task(**self._tasks[tid].task.model_dump()) for tid in page_task_ids]
# Determine next cursor
next_cursor = None
if start_index + self._page_size < len(all_task_ids) and page_task_ids:
next_cursor = page_task_ids[-1]
return tasks, next_cursor
async def delete_task(self, task_id: str) -> bool:
"""Delete a task."""
if task_id not in self._tasks:
return False
del self._tasks[task_id]
return True
async def wait_for_update(self, task_id: str) -> None:
"""Wait until the task status changes."""
if task_id not in self._tasks:
raise ValueError(f"Task with ID {task_id} not found")
# Create a fresh event for waiting (anyio.Event can't be cleared)
self._update_events[task_id] = anyio.Event()
event = self._update_events[task_id]
await event.wait()
async def notify_update(self, task_id: str) -> None:
"""Signal that a task has been updated."""
if task_id in self._update_events:
self._update_events[task_id].set()
# --- Testing/debugging helpers ---
def cleanup(self) -> None:
"""Cleanup all tasks (useful for testing or graceful shutdown)."""
self._tasks.clear()
self._update_events.clear()
def get_all_tasks(self) -> list[Task]:
"""Get all tasks (useful for debugging). Returns copies to prevent modification."""
self._cleanup_expired()
return [Task(**stored.task.model_dump()) for stored in self._tasks.values()]
@@ -0,0 +1,241 @@
"""
TaskMessageQueue - FIFO queue for task-related messages.
This implements the core message queue pattern from the MCP Tasks spec.
When a handler needs to send a request (like elicitation) during a task-augmented
request, the message is enqueued instead of sent directly. Messages are delivered
to the client only through the `tasks/result` endpoint.
This pattern enables:
1. Decoupling request handling from message delivery
2. Proper bidirectional communication via the tasks/result stream
3. Automatic status management (working <-> input_required)
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Literal
import anyio
from mcp.shared.experimental.tasks.resolver import Resolver
from mcp.types import JSONRPCNotification, JSONRPCRequest, RequestId
@dataclass
class QueuedMessage:
"""
A message queued for delivery via tasks/result.
Messages are stored with their type and a resolver for requests
that expect responses.
"""
type: Literal["request", "notification"]
"""Whether this is a request (expects response) or notification (one-way)."""
message: JSONRPCRequest | JSONRPCNotification
"""The JSON-RPC message to send."""
timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
"""When the message was enqueued."""
resolver: Resolver[dict[str, Any]] | None = None
"""Resolver to set when response arrives (only for requests)."""
original_request_id: RequestId | None = None
"""The original request ID used internally, for routing responses back."""
class TaskMessageQueue(ABC):
"""
Abstract interface for task message queuing.
This is a FIFO queue that stores messages to be delivered via `tasks/result`.
When a task-augmented handler calls elicit() or sends a notification, the
message is enqueued here instead of being sent directly to the client.
The `tasks/result` handler then dequeues and sends these messages through
the transport, with `relatedRequestId` set to the tasks/result request ID
so responses are routed correctly.
Implementations can use in-memory storage, Redis, etc.
"""
@abstractmethod
async def enqueue(self, task_id: str, message: QueuedMessage) -> None:
"""
Add a message to the queue for a task.
Args:
task_id: The task identifier
message: The message to enqueue
"""
@abstractmethod
async def dequeue(self, task_id: str) -> QueuedMessage | None:
"""
Remove and return the next message from the queue.
Args:
task_id: The task identifier
Returns:
The next message, or None if queue is empty
"""
@abstractmethod
async def peek(self, task_id: str) -> QueuedMessage | None:
"""
Return the next message without removing it.
Args:
task_id: The task identifier
Returns:
The next message, or None if queue is empty
"""
@abstractmethod
async def is_empty(self, task_id: str) -> bool:
"""
Check if the queue is empty for a task.
Args:
task_id: The task identifier
Returns:
True if no messages are queued
"""
@abstractmethod
async def clear(self, task_id: str) -> list[QueuedMessage]:
"""
Remove and return all messages from the queue.
This is useful for cleanup when a task is cancelled or completed.
Args:
task_id: The task identifier
Returns:
All queued messages (may be empty)
"""
@abstractmethod
async def wait_for_message(self, task_id: str) -> None:
"""
Wait until a message is available in the queue.
This blocks until either:
1. A message is enqueued for this task
2. The wait is cancelled
Args:
task_id: The task identifier
"""
@abstractmethod
async def notify_message_available(self, task_id: str) -> None:
"""
Signal that a message is available for a task.
This wakes up any coroutines waiting in wait_for_message().
Args:
task_id: The task identifier
"""
class InMemoryTaskMessageQueue(TaskMessageQueue):
"""
In-memory implementation of TaskMessageQueue.
This is suitable for single-process servers. For distributed systems,
implement TaskMessageQueue with Redis, RabbitMQ, etc.
Features:
- FIFO ordering per task
- Async wait for message availability
- Thread-safe for single-process async use
"""
def __init__(self) -> None:
self._queues: dict[str, list[QueuedMessage]] = {}
self._events: dict[str, anyio.Event] = {}
def _get_queue(self, task_id: str) -> list[QueuedMessage]:
"""Get or create the queue for a task."""
if task_id not in self._queues:
self._queues[task_id] = []
return self._queues[task_id]
async def enqueue(self, task_id: str, message: QueuedMessage) -> None:
"""Add a message to the queue."""
queue = self._get_queue(task_id)
queue.append(message)
# Signal that a message is available
await self.notify_message_available(task_id)
async def dequeue(self, task_id: str) -> QueuedMessage | None:
"""Remove and return the next message."""
queue = self._get_queue(task_id)
if not queue:
return None
return queue.pop(0)
async def peek(self, task_id: str) -> QueuedMessage | None:
"""Return the next message without removing it."""
queue = self._get_queue(task_id)
if not queue:
return None
return queue[0]
async def is_empty(self, task_id: str) -> bool:
"""Check if the queue is empty."""
queue = self._get_queue(task_id)
return len(queue) == 0
async def clear(self, task_id: str) -> list[QueuedMessage]:
"""Remove and return all messages."""
queue = self._get_queue(task_id)
messages = list(queue)
queue.clear()
return messages
async def wait_for_message(self, task_id: str) -> None:
"""Wait until a message is available."""
# Check if there are already messages
if not await self.is_empty(task_id):
return
# Create a fresh event for waiting (anyio.Event can't be cleared)
self._events[task_id] = anyio.Event()
event = self._events[task_id]
# Double-check after creating event (avoid race condition)
if not await self.is_empty(task_id):
return
# Wait for a new message
await event.wait()
async def notify_message_available(self, task_id: str) -> None:
"""Signal that a message is available."""
if task_id in self._events:
self._events[task_id].set()
def cleanup(self, task_id: str | None = None) -> None:
"""
Clean up queues and events.
Args:
task_id: If provided, clean up only this task. Otherwise clean up all.
"""
if task_id is not None:
self._queues.pop(task_id, None)
self._events.pop(task_id, None)
else:
self._queues.clear()
self._events.clear()
@@ -0,0 +1,45 @@
"""
Shared polling utilities for task operations.
This module provides generic polling logic that works for both client→server
and server→client task polling.
WARNING: These APIs are experimental and may change without notice.
"""
from collections.abc import AsyncIterator, Awaitable, Callable
import anyio
from mcp.shared.experimental.tasks.helpers import is_terminal
from mcp.types import GetTaskResult
async def poll_until_terminal(
get_task: Callable[[str], Awaitable[GetTaskResult]],
task_id: str,
default_interval_ms: int = 500,
) -> AsyncIterator[GetTaskResult]:
"""
Poll a task until it reaches terminal status.
This is a generic utility that works for both client→server and server→client
polling. The caller provides the get_task function appropriate for their direction.
Args:
get_task: Async function that takes task_id and returns GetTaskResult
task_id: The task to poll
default_interval_ms: Fallback poll interval if server doesn't specify
Yields:
GetTaskResult for each poll
"""
while True:
status = await get_task(task_id)
yield status
if is_terminal(status.status):
break
interval_ms = status.pollInterval if status.pollInterval is not None else default_interval_ms
await anyio.sleep(interval_ms / 1000)
@@ -0,0 +1,60 @@
"""
Resolver - An anyio-compatible future-like object for async result passing.
This provides a simple way to pass a result (or exception) from one coroutine
to another without depending on asyncio.Future.
"""
from typing import Generic, TypeVar, cast
import anyio
T = TypeVar("T")
class Resolver(Generic[T]):
"""
A simple resolver for passing results between coroutines.
Unlike asyncio.Future, this works with any anyio-compatible async backend.
Usage:
resolver: Resolver[str] = Resolver()
# In one coroutine:
resolver.set_result("hello")
# In another coroutine:
result = await resolver.wait() # returns "hello"
"""
def __init__(self) -> None:
self._event = anyio.Event()
self._value: T | None = None
self._exception: BaseException | None = None
def set_result(self, value: T) -> None:
"""Set the result value and wake up waiters."""
if self._event.is_set():
raise RuntimeError("Resolver already completed")
self._value = value
self._event.set()
def set_exception(self, exc: BaseException) -> None:
"""Set an exception and wake up waiters."""
if self._event.is_set():
raise RuntimeError("Resolver already completed")
self._exception = exc
self._event.set()
async def wait(self) -> T:
"""Wait for the result and return it, or raise the exception."""
await self._event.wait()
if self._exception is not None:
raise self._exception
# If we reach here, set_result() was called, so _value is set
return cast(T, self._value)
def done(self) -> bool:
"""Return True if the resolver has been completed."""
return self._event.is_set()
Loaded 100 of 1354 files, more files were not shown because too many files have changed in this diff. Show more