Use 120 characters instead of 88 (#856)

This commit is contained in:
Marcelo Trylesinski 2025-06-11 02:45:50 -07:00 • committed by GitHub
commit 543961968c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
45 changed files with 687 additions and 2142 deletions

View file

@ -21,9 +21,7 @@ def get_claude_config_path() -> Path | None:
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"
)
path = Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"), "Claude")
else:
return None
@ -37,8 +35,7 @@ def get_uv_path() -> str:
uv_path = shutil.which("uv")
if not uv_path:
logger.error(
"uv executable not found in PATH, falling back to 'uv'. "
"Please ensure uv is installed and in your PATH"
"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
@ -94,10 +91,7 @@ def update_claude_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]
):
if server_name in config["mcpServers"] and "env" in config["mcpServers"][server_name]:
existing_env = config["mcpServers"][server_name]["env"]
if env_vars:
# New vars take precedence over existing ones

View file

@ -45,9 +45,7 @@ def _get_npx_command():
# 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
)
subprocess.run([cmd, "--version"], check=True, capture_output=True, shell=True)
return cmd
except subprocess.CalledProcessError:
continue
@ -58,9 +56,7 @@ def _get_npx_command():
def _parse_env_var(env_var: str) -> tuple[str, str]:
"""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"
)
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()
@ -154,14 +150,10 @@ def _import_server(file: Path, server_object: str | None = None):
True if it's supported.
"""
if not isinstance(server_object, FastMCP):
logger.error(
f"The server object {object_name} is of type "
f"{type(server_object)} (expecting {FastMCP})."
)
logger.error(f"The server object {object_name} is of type " f"{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."
"Note that only FastMCP server is supported. Low level " "Server class is not yet supported."
)
return False
return True
@ -172,10 +164,7 @@ def _import_server(file: Path, server_object: str | None = None):
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"
)
logger.error(f"Ignoring object '{file}:{name}' as it's not a valid " "server object")
continue
return getattr(module, name)
@ -280,8 +269,7 @@ def dev(
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."
"npx not found. Please ensure Node.js and npm are properly installed " "and added to your system PATH."
)
sys.exit(1)
@ -383,8 +371,7 @@ def install(
typer.Option(
"--name",
"-n",
help="Custom name for the server (defaults to server's name attribute or"
" file name)",
help="Custom name for the server (defaults to server's name attribute or" " file name)",
),
] = None,
with_editable: Annotated[
@ -458,8 +445,7 @@ def install(
name = server.name
except (ImportError, ModuleNotFoundError) as e:
logger.debug(
"Could not import server (likely missing dependencies), using file"
" name",
"Could not import server (likely missing dependencies), using file" " name",
extra={"error": str(e)},
)
name = file.stem
@ -477,11 +463,7 @@ def install(
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
}
env_dict |= {k: v for k, v in dotenv.dotenv_values(env_file).items() if v is not None}
except Exception as e:
logger.error(f"Failed to load .env file: {e}")
sys.exit(1)

View file

@ -24,9 +24,7 @@ logger = logging.getLogger("client")
async def message_handler(
message: RequestResponder[types.ServerRequest, types.ClientResult]
| types.ServerNotification
| Exception,
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
) -> None:
if isinstance(message, Exception):
logger.error("Error: %s", message)
@ -60,9 +58,7 @@ async def main(command_or_url: str, args: list[str], env: list[tuple[str, str]])
await run_session(*streams)
else:
# Use stdio client for commands
server_parameters = StdioServerParameters(
command=command_or_url, args=args, env=env_dict
)
server_parameters = StdioServerParameters(command=command_or_url, args=args, env=env_dict)
async with stdio_client(server_parameters) as streams:
await run_session(*streams)

View file

@ -17,12 +17,7 @@ from urllib.parse import urlencode, urljoin
import anyio
import httpx
from mcp.shared.auth import (
OAuthClientInformationFull,
OAuthClientMetadata,
OAuthMetadata,
OAuthToken,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthMetadata, OAuthToken
from mcp.types import LATEST_PROTOCOL_VERSION
logger = logging.getLogger(__name__)
@ -100,10 +95,7 @@ class OAuthClientProvider(httpx.Auth):
def _generate_code_verifier(self) -> str:
"""Generate a cryptographically random code verifier for PKCE."""
return "".join(
secrets.choice(string.ascii_letters + string.digits + "-._~")
for _ in range(128)
)
return "".join(secrets.choice(string.ascii_letters + string.digits + "-._~") for _ in range(128))
def _generate_code_challenge(self, code_verifier: str) -> str:
"""Generate a code challenge from a code verifier using SHA256."""
@ -148,9 +140,7 @@ class OAuthClientProvider(httpx.Auth):
return None
response.raise_for_status()
metadata_json = response.json()
logger.debug(
f"OAuth metadata discovered (no MCP header): {metadata_json}"
)
logger.debug(f"OAuth metadata discovered (no MCP header): {metadata_json}")
return OAuthMetadata.model_validate(metadata_json)
except Exception:
logger.exception("Failed to discover OAuth metadata")
@ -176,17 +166,11 @@ class OAuthClientProvider(httpx.Auth):
registration_url = urljoin(auth_base_url, "/register")
# Handle default scope
if (
client_metadata.scope is None
and metadata
and metadata.scopes_supported is not None
):
if client_metadata.scope is None and metadata and metadata.scopes_supported is not None:
client_metadata.scope = " ".join(metadata.scopes_supported)
# Serialize client metadata
registration_data = client_metadata.model_dump(
by_alias=True, mode="json", exclude_none=True
)
registration_data = client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True)
async with httpx.AsyncClient() as client:
try:
@ -213,9 +197,7 @@ class OAuthClientProvider(httpx.Auth):
logger.exception("Registration error")
raise
async def async_auth_flow(
self, request: httpx.Request
) -> AsyncGenerator[httpx.Request, httpx.Response]:
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
"""
HTTPX auth flow integration.
"""
@ -225,9 +207,7 @@ class OAuthClientProvider(httpx.Auth):
await self.ensure_token()
# Add Bearer token if available
if self._current_tokens and self._current_tokens.access_token:
request.headers["Authorization"] = (
f"Bearer {self._current_tokens.access_token}"
)
request.headers["Authorization"] = f"Bearer {self._current_tokens.access_token}"
response = yield request
@ -305,11 +285,7 @@ class OAuthClientProvider(httpx.Auth):
return
# Try refreshing existing token
if (
self._current_tokens
and self._current_tokens.refresh_token
and await self._refresh_access_token()
):
if self._current_tokens and self._current_tokens.refresh_token and await self._refresh_access_token():
return
# Fall back to full OAuth flow
@ -361,12 +337,8 @@ class OAuthClientProvider(httpx.Auth):
auth_code, returned_state = await self.callback_handler()
# Validate state parameter for CSRF protection
if returned_state is None or not secrets.compare_digest(
returned_state, self._auth_state
):
raise Exception(
f"State parameter mismatch: {returned_state} != {self._auth_state}"
)
if returned_state is None or not secrets.compare_digest(returned_state, self._auth_state):
raise Exception(f"State parameter mismatch: {returned_state} != {self._auth_state}")
# Clear state after validation
self._auth_state = None
@ -377,9 +349,7 @@ class OAuthClientProvider(httpx.Auth):
# Exchange authorization code for tokens
await self._exchange_code_for_token(auth_code, client_info)
async def _exchange_code_for_token(
self, auth_code: str, client_info: OAuthClientInformationFull
) -> None:
async def _exchange_code_for_token(self, auth_code: str, client_info: OAuthClientInformationFull) -> None:
"""Exchange authorization code for access token."""
# Get token endpoint
if self._metadata and self._metadata.token_endpoint:
@ -412,17 +382,10 @@ class OAuthClientProvider(httpx.Auth):
# Parse OAuth error response
try:
error_data = response.json()
error_msg = error_data.get(
"error_description", error_data.get("error", "Unknown error")
)
raise Exception(
f"Token exchange failed: {error_msg} "
f"(HTTP {response.status_code})"
)
error_msg = error_data.get("error_description", error_data.get("error", "Unknown error"))
raise Exception(f"Token exchange failed: {error_msg} " f"(HTTP {response.status_code})")
except Exception:
raise Exception(
f"Token exchange failed: {response.status_code} {response.text}"
)
raise Exception(f"Token exchange failed: {response.status_code} {response.text}")
# Parse token response
token_response = OAuthToken.model_validate(response.json())

View file

@ -38,16 +38,12 @@ class LoggingFnT(Protocol):
class MessageHandlerFnT(Protocol):
async def __call__(
self,
message: RequestResponder[types.ServerRequest, types.ClientResult]
| types.ServerNotification
| Exception,
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
) -> None: ...
async def _default_message_handler(
message: RequestResponder[types.ServerRequest, types.ClientResult]
| types.ServerNotification
| Exception,
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
) -> None:
await anyio.lowlevel.checkpoint()
@ -77,9 +73,7 @@ async def _default_logging_callback(
pass
ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(
types.ClientResult | types.ErrorData
)
ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData)
class ClientSession(
@ -116,11 +110,7 @@ class ClientSession(
self._message_handler = message_handler or _default_message_handler
async def initialize(self) -> types.InitializeResult:
sampling = (
types.SamplingCapability()
if self._sampling_callback is not _default_sampling_callback
else None
)
sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
roots = (
# TODO: Should this be based on whether we
# _will_ send notifications, or only whether
@ -149,15 +139,10 @@ class ClientSession(
)
if result.protocolVersion not in SUPPORTED_PROTOCOL_VERSIONS:
raise RuntimeError(
"Unsupported protocol version from the server: "
f"{result.protocolVersion}"
)
raise RuntimeError("Unsupported protocol version from the server: " f"{result.protocolVersion}")
await self.send_notification(
types.ClientNotification(
types.InitializedNotification(method="notifications/initialized")
)
types.ClientNotification(types.InitializedNotification(method="notifications/initialized"))
)
return result
@ -207,33 +192,25 @@ class ClientSession(
types.EmptyResult,
)
async def list_resources(
self, cursor: str | None = None
) -> types.ListResourcesResult:
async def list_resources(self, cursor: str | None = None) -> types.ListResourcesResult:
"""Send a resources/list request."""
return await self.send_request(
types.ClientRequest(
types.ListResourcesRequest(
method="resources/list",
params=types.PaginatedRequestParams(cursor=cursor)
if cursor is not None
else None,
params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None,
)
),
types.ListResourcesResult,
)
async def list_resource_templates(
self, cursor: str | None = None
) -> types.ListResourceTemplatesResult:
async def list_resource_templates(self, cursor: str | None = None) -> types.ListResourceTemplatesResult:
"""Send a resources/templates/list request."""
return await self.send_request(
types.ClientRequest(
types.ListResourceTemplatesRequest(
method="resources/templates/list",
params=types.PaginatedRequestParams(cursor=cursor)
if cursor is not None
else None,
params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None,
)
),
types.ListResourceTemplatesResult,
@ -305,17 +282,13 @@ class ClientSession(
types.ClientRequest(
types.ListPromptsRequest(
method="prompts/list",
params=types.PaginatedRequestParams(cursor=cursor)
if cursor is not None
else None,
params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None,
)
),
types.ListPromptsResult,
)
async def get_prompt(
self, name: str, arguments: dict[str, str] | None = None
) -> types.GetPromptResult:
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(
@ -352,9 +325,7 @@ class ClientSession(
types.ClientRequest(
types.ListToolsRequest(
method="tools/list",
params=types.PaginatedRequestParams(cursor=cursor)
if cursor is not None
else None,
params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None,
)
),
types.ListToolsResult,
@ -370,9 +341,7 @@ class ClientSession(
)
)
async def _received_request(
self, responder: RequestResponder[types.ServerRequest, types.ClientResult]
) -> None:
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,
@ -395,22 +364,16 @@ class ClientSession(
case types.PingRequest():
with responder:
return await responder.respond(
types.ClientResult(root=types.EmptyResult())
)
return await responder.respond(types.ClientResult(root=types.EmptyResult()))
async def _handle_incoming(
self,
req: RequestResponder[types.ServerRequest, types.ClientResult]
| types.ServerNotification
| Exception,
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:
async def _received_notification(self, notification: types.ServerNotification) -> None:
"""Handle notifications from the server."""
# Process specific notification types
match notification.root:

View file

@ -62,9 +62,7 @@ class StreamableHttpParameters(BaseModel):
terminate_on_close: bool = True
ServerParameters: TypeAlias = (
StdioServerParameters | SseServerParameters | StreamableHttpParameters
)
ServerParameters: TypeAlias = StdioServerParameters | SseServerParameters | StreamableHttpParameters
class ClientSessionGroup:
@ -261,9 +259,7 @@ class ClientSessionGroup:
)
read, write, _ = await session_stack.enter_async_context(client)
session = await session_stack.enter_async_context(
mcp.ClientSession(read, write)
)
session = await session_stack.enter_async_context(mcp.ClientSession(read, write))
result = await session.initialize()
# Session successfully initialized.
@ -280,9 +276,7 @@ class ClientSessionGroup:
await session_stack.aclose()
raise
async def _aggregate_components(
self, server_info: types.Implementation, session: mcp.ClientSession
) -> None:
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

View file

@ -73,20 +73,16 @@ async def sse_client(
match sse.event:
case "endpoint":
endpoint_url = urljoin(url, sse.data)
logger.debug(
f"Received endpoint URL: {endpoint_url}"
)
logger.debug(f"Received endpoint URL: {endpoint_url}")
url_parsed = urlparse(url)
endpoint_parsed = urlparse(endpoint_url)
if (
url_parsed.netloc != endpoint_parsed.netloc
or url_parsed.scheme
!= endpoint_parsed.scheme
or url_parsed.scheme != endpoint_parsed.scheme
):
error_msg = (
"Endpoint origin does not match "
f"connection origin: {endpoint_url}"
"Endpoint origin does not match " f"connection origin: {endpoint_url}"
)
logger.error(error_msg)
raise ValueError(error_msg)
@ -98,22 +94,16 @@ async def sse_client(
message = types.JSONRPCMessage.model_validate_json( # noqa: E501
sse.data
)
logger.debug(
f"Received server message: {message}"
)
logger.debug(f"Received server message: {message}")
except Exception as exc:
logger.error(
f"Error parsing server message: {exc}"
)
logger.error(f"Error parsing server message: {exc}")
await read_stream_writer.send(exc)
continue
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
case _:
logger.warning(
f"Unknown SSE event: {sse.event}"
)
logger.warning(f"Unknown SSE event: {sse.event}")
except Exception as exc:
logger.error(f"Error in sse_reader: {exc}")
await read_stream_writer.send(exc)
@ -124,9 +114,7 @@ async def sse_client(
try:
async with write_stream_reader:
async for session_message in write_stream_reader:
logger.debug(
f"Sending client message: {session_message}"
)
logger.debug(f"Sending client message: {session_message}")
response = await client.post(
endpoint_url,
json=session_message.message.model_dump(
@ -136,19 +124,14 @@ async def sse_client(
),
)
response.raise_for_status()
logger.debug(
"Client message sent successfully: "
f"{response.status_code}"
)
logger.debug("Client message sent successfully: " f"{response.status_code}")
except Exception as exc:
logger.error(f"Error in post_writer: {exc}")
finally:
await write_stream.aclose()
endpoint_url = await tg.start(sse_reader)
logger.debug(
f"Starting post writer with endpoint URL: {endpoint_url}"
)
logger.debug(f"Starting post writer with endpoint URL: {endpoint_url}")
tg.start_soon(post_writer, endpoint_url)
try:

View file

@ -115,11 +115,7 @@ async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stder
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()
),
env=({**get_default_environment(), **server.env} if server.env is not None else get_default_environment()),
errlog=errlog,
cwd=server.cwd,
)
@ -163,9 +159,7 @@ async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stder
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
)
json = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
await process.stdin.send(
(json + "\n").encode(
encoding=server.encoding,
@ -229,8 +223,6 @@ async def _create_platform_compatible_process(
if sys.platform == "win32":
process = await create_windows_process(command, args, env, errlog, cwd)
else:
process = await anyio.open_process(
[command, *args], env=env, stderr=errlog, cwd=cwd
)
process = await anyio.open_process([command, *args], env=env, stderr=errlog, cwd=cwd)
return process

View file

@ -82,9 +82,7 @@ async def create_windows_process(
return process
except Exception:
# Don't raise, let's try to create the process without creation flags
process = await anyio.open_process(
[command, *args], env=env, stderr=errlog, cwd=cwd
)
process = await anyio.open_process([command, *args], env=env, stderr=errlog, cwd=cwd)
return process

View file

@ -106,9 +106,7 @@ class StreamableHTTPTransport:
**self.headers,
}
def _update_headers_with_session(
self, base_headers: dict[str, str]
) -> dict[str, str]:
def _update_headers_with_session(self, base_headers: dict[str, str]) -> dict[str, str]:
"""Update headers with session ID if available."""
headers = base_headers.copy()
if self.session_id:
@ -117,17 +115,11 @@ class StreamableHTTPTransport:
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"
)
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"
)
return isinstance(message.root, JSONRPCNotification) and message.root.method == "notifications/initialized"
def _maybe_extract_session_id_from_response(
self,
@ -153,9 +145,7 @@ class StreamableHTTPTransport:
logger.debug(f"SSE 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
):
if original_request_id is not None and isinstance(message.root, JSONRPCResponse | JSONRPCError):
message.root.id = original_request_id
session_message = SessionMessage(message)
@ -227,7 +217,8 @@ class StreamableHTTPTransport:
self.url,
headers=headers,
timeout=httpx.Timeout(
self.timeout.total_seconds(), read=ctx.sse_read_timeout.total_seconds()
self.timeout.total_seconds(),
read=ctx.sse_read_timeout.total_seconds(),
),
) as event_source:
event_source.response.raise_for_status()
@ -298,9 +289,7 @@ class StreamableHTTPTransport:
logger.error(f"Error parsing JSON response: {exc}")
await read_stream_writer.send(exc)
async def _handle_sse_response(
self, response: httpx.Response, ctx: RequestContext
) -> None:
async def _handle_sse_response(self, response: httpx.Response, ctx: RequestContext) -> None:
"""Handle SSE response from the server."""
try:
event_source = EventSource(response)
@ -308,11 +297,7 @@ class StreamableHTTPTransport:
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
),
resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None),
)
# If the SSE event indicates completion, like returning respose/error
# break the loop
@ -455,12 +440,8 @@ async def streamablehttp_client(
"""
transport = StreamableHTTPTransport(url, headers, timeout, sse_read_timeout, auth)
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)
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)
async with anyio.create_task_group() as tg:
try:
@ -476,9 +457,7 @@ async def streamablehttp_client(
) as client:
# Define callbacks that need access to tg
def start_get_stream() -> None:
tg.start_soon(
transport.handle_get_stream, client, read_stream_writer
)
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
tg.start_soon(
transport.post_writer,

View file

@ -19,10 +19,7 @@ logger = logging.getLogger(__name__)
async def websocket_client(
url: str,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
],
tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]],
None,
]:
"""
@ -74,9 +71,7 @@ async def websocket_client(
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
)
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:

View file

@ -2,7 +2,4 @@ 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()
)
return "\n".join(f"{'.'.join(str(loc) for loc in e['loc'])}: {e['msg']}" for e in validation_error.errors())

View file

@ -7,9 +7,7 @@ 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.errors import stringify_pydantic_error
from mcp.server.auth.json_response import PydanticJSONResponse
from mcp.server.auth.provider import (
AuthorizationErrorCode,
@ -18,10 +16,7 @@ from mcp.server.auth.provider import (
OAuthAuthorizationServerProvider,
construct_redirect_uri,
)
from mcp.shared.auth import (
InvalidRedirectUriError,
InvalidScopeError,
)
from mcp.shared.auth import InvalidRedirectUriError, InvalidScopeError
logger = logging.getLogger(__name__)
@ -29,23 +24,16 @@ 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"
)
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"
)
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"
)
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",
description="Optional scope; if specified, should be " "a space-separated list of scope strings",
)
@ -57,9 +45,7 @@ class AuthorizationErrorResponse(BaseModel):
state: str | None = None
def best_effort_extract_string(
key: str, params: None | FormData | QueryParams
) -> str | None:
def best_effort_extract_string(key: str, params: None | FormData | QueryParams) -> str | None:
if params is None:
return None
value = params.get(key)
@ -138,9 +124,7 @@ class AuthorizationHandler:
if redirect_uri and client:
return RedirectResponse(
url=construct_redirect_uri(
str(redirect_uri), **error_resp.model_dump(exclude_none=True)
),
url=construct_redirect_uri(str(redirect_uri), **error_resp.model_dump(exclude_none=True)),
status_code=302,
headers={"Cache-Control": "no-store"},
)
@ -172,9 +156,7 @@ class AuthorizationHandler:
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)
)
return await error_response(error, stringify_pydantic_error(validation_error))
# Get client information
client = await self.provider.get_client(
@ -229,16 +211,9 @@ class AuthorizationHandler:
)
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,
)
return await error_response(error=e.error, error_description=e.error_description)
except Exception as validation_error:
# 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"
)
logger.exception("Unexpected error in authorization_handler", exc_info=validation_error)
return await error_response(error="server_error", error_description="An unexpected error occurred")

View file

@ -10,11 +10,7 @@ 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.provider import OAuthAuthorizationServerProvider, RegistrationError, RegistrationErrorCode
from mcp.server.auth.settings import ClientRegistrationOptions
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
@ -60,9 +56,7 @@ class RegistrationHandler:
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
):
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):
@ -78,8 +72,7 @@ class RegistrationHandler:
return PydanticJSONResponse(
content=RegistrationErrorResponse(
error="invalid_client_metadata",
error_description="grant_types must be authorization_code "
"and refresh_token",
error_description="grant_types must be authorization_code " "and refresh_token",
),
status_code=400,
)
@ -122,8 +115,6 @@ class RegistrationHandler:
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
),
content=RegistrationErrorResponse(error=e.error, error_description=e.error_description),
status_code=400,
)

View file

@ -10,15 +10,8 @@ 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,
)
from mcp.server.auth.middleware.client_auth import AuthenticationError, ClientAuthenticator
from mcp.server.auth.provider import AccessToken, OAuthAuthorizationServerProvider, RefreshToken
class RevocationRequest(BaseModel):

View file

@ -7,19 +7,10 @@ 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.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.server.auth.middleware.client_auth import AuthenticationError, ClientAuthenticator
from mcp.server.auth.provider import OAuthAuthorizationServerProvider, TokenError, TokenErrorCode
from mcp.shared.auth import OAuthToken
@ -27,9 +18,7 @@ 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"
)
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
@ -127,8 +116,7 @@ class TokenHandler:
TokenErrorResponse(
error="unsupported_grant_type",
error_description=(
f"Unsupported grant type (supported grant types are "
f"{client_info.grant_types})"
f"Unsupported grant type (supported grant types are " f"{client_info.grant_types})"
),
)
)
@ -137,9 +125,7 @@ class TokenHandler:
match token_request:
case AuthorizationCodeRequest():
auth_code = await self.provider.load_authorization_code(
client_info, token_request.code
)
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(
@ -169,18 +155,13 @@ class TokenHandler:
return self.response(
TokenErrorResponse(
error="invalid_request",
error_description=(
"redirect_uri did not match the one "
"used when creating auth code"
),
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("=")
)
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
@ -193,9 +174,7 @@ class TokenHandler:
try:
# Exchange authorization code for tokens
tokens = await self.provider.exchange_authorization_code(
client_info, auth_code
)
tokens = await self.provider.exchange_authorization_code(client_info, auth_code)
except TokenError as e:
return self.response(
TokenErrorResponse(
@ -205,13 +184,8 @@ class TokenHandler:
)
case RefreshTokenRequest():
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
):
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(
@ -230,29 +204,20 @@ class TokenHandler:
)
# Parse scopes if provided
scopes = (
token_request.scope.split(" ")
if token_request.scope
else refresh_token.scopes
)
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"
),
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
)
tokens = await self.provider.exchange_refresh_token(client_info, refresh_token, scopes)
except TokenError as e:
return self.response(
TokenErrorResponse(

View file

@ -7,9 +7,7 @@ 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
)
auth_context_var = contextvars.ContextVar[AuthenticatedUser | None]("auth_context", default=None)
def get_access_token() -> AccessToken | None:

View file

@ -1,11 +1,7 @@
import time
from typing import Any
from starlette.authentication import (
AuthCredentials,
AuthenticationBackend,
SimpleUser,
)
from starlette.authentication import AuthCredentials, AuthenticationBackend, SimpleUser
from starlette.exceptions import HTTPException
from starlette.requests import HTTPConnection
from starlette.types import Receive, Scope, Send
@ -35,11 +31,7 @@ class BearerAuthBackend(AuthenticationBackend):
async def authenticate(self, conn: HTTPConnection):
auth_header = next(
(
conn.headers.get(key)
for key in conn.headers
if key.lower() == "authorization"
),
(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 "):
@ -87,10 +79,7 @@ class RequireAuthMiddleware:
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
):
if auth_credentials is None or required_scope not in auth_credentials.scopes:
raise HTTPException(status_code=403, detail="Insufficient scope")
await self.app(scope, receive, send)

View file

@ -30,9 +30,7 @@ class ClientAuthenticator:
"""
self.provider = provider
async def authenticate(
self, client_id: str, client_secret: str | None
) -> OAuthClientInformationFull:
async def authenticate(self, client_id: str, client_secret: str | None) -> OAuthClientInformationFull:
# Look up client information
client = await self.provider.get_client(client_id)
if not client:
@ -47,10 +45,7 @@ class ClientAuthenticator:
if client.client_secret != client_secret:
raise AuthenticationError("Invalid client_secret")
if (
client.client_secret_expires_at
and client.client_secret_expires_at < int(time.time())
):
if client.client_secret_expires_at and client.client_secret_expires_at < int(time.time()):
raise AuthenticationError("Client secret has expired")
return client

View file

@ -4,10 +4,7 @@ from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
from pydantic import AnyUrl, BaseModel
from mcp.shared.auth import (
OAuthClientInformationFull,
OAuthToken,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
class AuthorizationParams(BaseModel):
@ -96,9 +93,7 @@ RefreshTokenT = TypeVar("RefreshTokenT", bound=RefreshToken)
AccessTokenT = TypeVar("AccessTokenT", bound=AccessToken)
class OAuthAuthorizationServerProvider(
Protocol, Generic[AuthorizationCodeT, RefreshTokenT, AccessTokenT]
):
class OAuthAuthorizationServerProvider(Protocol, Generic[AuthorizationCodeT, RefreshTokenT, AccessTokenT]):
async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
"""
Retrieves client information by client ID.
@ -129,9 +124,7 @@ class OAuthAuthorizationServerProvider(
"""
...
async def authorize(
self, client: OAuthClientInformationFull, params: AuthorizationParams
) -> str:
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.
@ -207,9 +200,7 @@ class OAuthAuthorizationServerProvider(
"""
...
async def load_refresh_token(
self, client: OAuthClientInformationFull, refresh_token: str
) -> RefreshTokenT | None:
async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshTokenT | None:
"""
Loads a RefreshToken by its token string.

Some files were not shown because too many files have changed in this diff Show more