FrameBridge 0.3.1 community source release
This commit is contained in:
commit
b23819fff1
1354 files changed
+343716
No files matched your search
+135
@@ -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",
|
||||
]
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
"""FastMCP CLI package."""
|
||||
|
||||
from .cli import app
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
app()
|
||||
+148
@@ -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.
+85
@@ -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()
|
||||
+21
@@ -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",
|
||||
]
|
||||
+10
@@ -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."""
|
||||
Vendored
Whitespace-only changes.
+487
@@ -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"}
|
||||
)
|
||||
+616
@@ -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
|
||||
+336
@@ -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}")
|
||||
Vendored
+9
@@ -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"]
|
||||
+290
@@ -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
|
||||
Vendored
+224
@@ -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
|
||||
+615
@@ -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
|
||||
+447
@@ -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
|
||||
+164
@@ -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()
|
||||
+278
@@ -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)
|
||||
+722
@@ -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
|
||||
+86
@@ -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()
|
||||
+1
@@ -0,0 +1 @@
|
||||
"""Platform-specific utilities for MCP."""
|
||||
+1
@@ -0,0 +1 @@
|
||||
"""POSIX-specific utilities for MCP."""
|
||||
+60
@@ -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}")
|
||||
+1
@@ -0,0 +1 @@
|
||||
"""Windows-specific utilities for MCP."""
|
||||
+338
@@ -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.
+5
@@ -0,0 +1,5 @@
|
||||
from .fastmcp import FastMCP
|
||||
from .lowlevel import NotificationOptions, Server
|
||||
from .models import InitializationOptions
|
||||
|
||||
__all__ = ["Server", "FastMCP", "NotificationOptions", "InitializationOptions"]
|
||||
+50
@@ -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")
|
||||
+3
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
MCP OAuth server authorization components.
|
||||
"""
|
||||
+5
@@ -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())
|
||||
Vendored
+3
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
Request handlers for MCP authorization endpoints.
|
||||
"""
|
||||
Vendored
+224
@@ -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")
|
||||
Vendored
+29
@@ -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
|
||||
)
|
||||
Vendored
+136
@@ -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,
|
||||
)
|
||||
Vendored
+91
@@ -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",
|
||||
},
|
||||
)
|
||||
Vendored
+241
@@ -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))
|
||||
Vendored
+10
@@ -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")
|
||||
Vendored
+3
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
Middleware for MCP authorization.
|
||||
"""
|
||||
+48
@@ -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)
|
||||
+128
@@ -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,
|
||||
}
|
||||
)
|
||||
+115
@@ -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
|
||||
+301
@@ -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)
|
||||
+253
@@ -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"],
|
||||
)
|
||||
]
|
||||
+30
@@ -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.",
|
||||
)
|
||||
+190
@@ -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}")
|
||||
Vendored
+11
@@ -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
|
||||
"""
|
||||
+238
@@ -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 {})
|
||||
+220
@@ -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)
|
||||
source/sources/original-geranas/FrameBridge-Fleet-Hub/vendor/mcp/server/experimental/task_context.py
Vendored
+612
@@ -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
|
||||
+235
@@ -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
|
||||
source/sources/original-geranas/FrameBridge-Fleet-Hub/vendor/mcp/server/experimental/task_support.py
Vendored
+115
@@ -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(),
|
||||
)
|
||||
Vendored
+11
@@ -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"]
|
||||
Vendored
+21
@@ -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."""
|
||||
Vendored
+4
@@ -0,0 +1,4 @@
|
||||
from .base import Prompt
|
||||
from .manager import PromptManager
|
||||
|
||||
__all__ = ["Prompt", "PromptManager"]
|
||||
Vendored
+183
@@ -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}")
|
||||
Vendored
+60
@@ -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)
|
||||
+23
@@ -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",
|
||||
]
|
||||
Vendored
+50
@@ -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
|
||||
+115
@@ -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())
|
||||
+122
@@ -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}")
|
||||
Vendored
+203
@@ -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}")
|
||||
+1349
File diff suppressed because it is too large.
Load diff
Vendored
+4
@@ -0,0 +1,4 @@
|
||||
from .base import Tool
|
||||
from .tool_manager import ToolManager
|
||||
|
||||
__all__ = ["Tool", "ToolManager"]
|
||||
Vendored
+126
@@ -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))
|
||||
)
|
||||
+93
@@ -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)
|
||||
+1
@@ -0,0 +1 @@
|
||||
"""FastMCP utility modules."""
|
||||
+68
@@ -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
|
||||
+533
@@ -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)]
|
||||
source/sources/original-geranas/FrameBridge-Fleet-Hub/vendor/mcp/server/fastmcp/utilities/logging.py
Vendored
+43
@@ -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,
|
||||
)
|
||||
Vendored
+101
@@ -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)
|
||||
Vendored
+3
@@ -0,0 +1,3 @@
|
||||
from .server import NotificationOptions, Server
|
||||
|
||||
__all__ = ["Server", "NotificationOptions"]
|
||||
Vendored
+288
@@ -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
|
||||
Vendored
+54
@@ -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()
|
||||
Vendored
+11
@@ -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
|
||||
+803
@@ -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())
|
||||
+20
@@ -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
|
||||
+691
@@ -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
|
||||
+249
@@ -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)
|
||||
+88
@@ -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
|
||||
+1067
File diff suppressed because it is too large.
Load diff
Vendored
+297
@@ -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)
|
||||
Vendored
+127
@@ -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
|
||||
+104
@@ -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")
|
||||
+62
@@ -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.
+87
@@ -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)
|
||||
+159
@@ -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
|
||||
+85
@@ -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)
|
||||
+32
@@ -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
|
||||
+71
@@ -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)
|
||||
Vendored
+7
@@ -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.
|
||||
"""
|
||||
+12
@@ -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
|
||||
"""
|
||||
+115
@@ -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",
|
||||
)
|
||||
)
|
||||
+101
@@ -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,
|
||||
)
|
||||
+181
@@ -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
|
||||
+219
@@ -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()]
|
||||
+241
@@ -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()
|
||||
+45
@@ -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)
|
||||
+60
@@ -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
Reference in new issue
Block a user