Skip to content

Commit 6f0e8e7

Browse files
Kludexclaude
andauthored
refactor: code style improvements and formatting cleanup (#1962)
Co-authored-by: Claude Opus 4.5 <[email protected]>
1 parent 1b5287c commit 6f0e8e7

File tree

7 files changed

+30
-91
lines changed

7 files changed

+30
-91
lines changed

src/mcp/client/stdio.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,10 +177,7 @@ async def stdin_writer():
177177
except anyio.ClosedResourceError: # pragma: no cover
178178
await anyio.lowlevel.checkpoint()
179179

180-
async with (
181-
anyio.create_task_group() as tg,
182-
process,
183-
):
180+
async with anyio.create_task_group() as tg, process:
184181
tg.start_soon(stdout_reader)
185182
tg.start_soon(stdin_writer)
186183
try:

src/mcp/server/lowlevel/server.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -825,7 +825,7 @@ def streamable_http_app(
825825
host: str = "127.0.0.1",
826826
auth: AuthSettings | None = None,
827827
token_verifier: TokenVerifier | None = None,
828-
auth_server_provider: (OAuthAuthorizationServerProvider[Any, Any, Any] | None) = None,
828+
auth_server_provider: OAuthAuthorizationServerProvider[Any, Any, Any] | None = None,
829829
custom_starlette_routes: list[Route] | None = None,
830830
debug: bool = False,
831831
) -> Starlette:

src/mcp/server/models.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,7 @@
44

55
from pydantic import BaseModel
66

7-
from mcp.types import (
8-
Icon,
9-
ServerCapabilities,
10-
)
7+
from mcp.types import Icon, ServerCapabilities
118

129

1310
class InitializationOptions(BaseModel):

src/mcp/server/stdio.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,7 @@ async def run_server():
3030

3131

3232
@asynccontextmanager
33-
async def stdio_server(
34-
stdin: anyio.AsyncFile[str] | None = None,
35-
stdout: anyio.AsyncFile[str] | None = None,
36-
):
33+
async def stdio_server(stdin: anyio.AsyncFile[str] | None = None, stdout: anyio.AsyncFile[str] | None = None):
3734
"""Server transport for stdio: this communicates with an MCP client by reading
3835
from the current process' stdin and writing to stdout.
3936
"""

src/mcp/server/streamable_http_manager.py

Lines changed: 8 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -124,20 +124,10 @@ async def lifespan(app: Starlette) -> AsyncIterator[None]:
124124
# Clear any remaining server instances
125125
self._server_instances.clear()
126126

127-
async def handle_request(
128-
self,
129-
scope: Scope,
130-
receive: Receive,
131-
send: Send,
132-
) -> None:
127+
async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
133128
"""Process ASGI request with proper session handling and transport setup.
134129
135130
Dispatches to the appropriate handler based on stateless mode.
136-
137-
Args:
138-
scope: ASGI scope
139-
receive: ASGI receive function
140-
send: ASGI send function
141131
"""
142132
if self._task_group is None:
143133
raise RuntimeError("Task group is not initialized. Make sure to use run().")
@@ -148,19 +138,8 @@ async def handle_request(
148138
else:
149139
await self._handle_stateful_request(scope, receive, send)
150140

151-
async def _handle_stateless_request(
152-
self,
153-
scope: Scope,
154-
receive: Receive,
155-
send: Send,
156-
) -> None:
157-
"""Process request in stateless mode - creating a new transport for each request.
158-
159-
Args:
160-
scope: ASGI scope
161-
receive: ASGI receive function
162-
send: ASGI send function
163-
"""
141+
async def _handle_stateless_request(self, scope: Scope, receive: Receive, send: Send) -> None:
142+
"""Process request in stateless mode - creating a new transport for each request."""
164143
logger.debug("Stateless mode: Creating new transport for this request")
165144
# No session ID needed in stateless mode
166145
http_transport = StreamableHTTPServerTransport(
@@ -196,19 +175,8 @@ async def run_stateless_server(*, task_status: TaskStatus[None] = anyio.TASK_STA
196175
# Terminate the transport after the request is handled
197176
await http_transport.terminate()
198177

199-
async def _handle_stateful_request(
200-
self,
201-
scope: Scope,
202-
receive: Receive,
203-
send: Send,
204-
) -> None:
205-
"""Process request in stateful mode - maintaining session state between requests.
206-
207-
Args:
208-
scope: ASGI scope
209-
receive: ASGI receive function
210-
send: ASGI send function
211-
"""
178+
async def _handle_stateful_request(self, scope: Scope, receive: Receive, send: Send) -> None:
179+
"""Process request in stateful mode - maintaining session state between requests."""
212180
request = Request(scope, receive)
213181
request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
214182

@@ -248,11 +216,8 @@ async def run_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORE
248216
self.app.create_initialization_options(),
249217
stateless=False, # Stateful mode
250218
)
251-
except Exception as e:
252-
logger.error(
253-
f"Session {http_transport.mcp_session_id} crashed: {e}",
254-
exc_info=True,
255-
)
219+
except Exception:
220+
logger.exception(f"Session {http_transport.mcp_session_id} crashed")
256221
finally:
257222
# Only remove from instances if not terminated
258223
if ( # pragma: no branch
@@ -262,8 +227,7 @@ async def run_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORE
262227
):
263228
logger.info(
264229
"Cleaning up crashed session "
265-
f"{http_transport.mcp_session_id} from "
266-
"active instances."
230+
f"{http_transport.mcp_session_id} from active instances."
267231
)
268232
del self._server_instances[http_transport.mcp_session_id]
269233

src/mcp/server/transport_security.py

Lines changed: 17 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -9,37 +9,35 @@
99
logger = logging.getLogger(__name__)
1010

1111

12+
# TODO(Marcelo): We should flatten these settings. To be fair, I don't think we should even have this middleware.
1213
class TransportSecuritySettings(BaseModel):
1314
"""Settings for MCP transport security features.
1415
15-
These settings help protect against DNS rebinding attacks by validating
16-
incoming request headers.
16+
These settings help protect against DNS rebinding attacks by validating incoming request headers.
1717
"""
1818

19-
enable_dns_rebinding_protection: bool = Field(
20-
default=True,
21-
description="Enable DNS rebinding protection (recommended for production)",
22-
)
19+
enable_dns_rebinding_protection: bool = True
20+
"""Enable DNS rebinding protection (recommended for production)."""
2321

24-
allowed_hosts: list[str] = Field(
25-
default=[],
26-
description="List of allowed Host header values. Only applies when "
27-
+ "enable_dns_rebinding_protection is True.",
28-
)
22+
allowed_hosts: list[str] = Field(default_factory=list)
23+
"""List of allowed Host header values.
2924
30-
allowed_origins: list[str] = Field(
31-
default=[],
32-
description="List of allowed Origin header values. Only applies when "
33-
+ "enable_dns_rebinding_protection is True.",
34-
)
25+
Only applies when `enable_dns_rebinding_protection` is `True`.
26+
"""
27+
28+
allowed_origins: list[str] = Field(default_factory=list)
29+
"""List of allowed Origin header values.
30+
31+
Only applies when `enable_dns_rebinding_protection` is `True`.
32+
"""
3533

3634

35+
# TODO(Marcelo): This should be a proper ASGI middleware. I'm sad to see this.
3736
class TransportSecurityMiddleware:
3837
"""Middleware to enforce DNS rebinding protection for MCP transport endpoints."""
3938

4039
def __init__(self, settings: TransportSecuritySettings | None = None):
41-
# If not specified, disable DNS rebinding protection by default
42-
# for backwards compatibility
40+
# If not specified, disable DNS rebinding protection by default for backwards compatibility
4341
self.settings = settings or TransportSecuritySettings(enable_dns_rebinding_protection=False)
4442

4543
def _validate_host(self, host: str | None) -> bool: # pragma: no cover
@@ -88,16 +86,7 @@ def _validate_origin(self, origin: str | None) -> bool: # pragma: no cover
8886

8987
def _validate_content_type(self, content_type: str | None) -> bool:
9088
"""Validate the Content-Type header for POST requests."""
91-
if not content_type: # pragma: lax no cover
92-
logger.warning("Missing Content-Type header in POST request")
93-
return False
94-
95-
# Content-Type must start with application/json
96-
if not content_type.lower().startswith("application/json"):
97-
logger.warning(f"Invalid Content-Type header: {content_type}")
98-
return False
99-
100-
return True
89+
return content_type is not None and content_type.lower().startswith("application/json")
10190

10291
async def validate_request(self, request: Request, is_post: bool = False) -> Response | None:
10392
"""Validate request headers for DNS rebinding protection.

src/mcp/shared/memory.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,5 @@ async def create_client_server_memory_streams() -> AsyncGenerator[tuple[MessageS
2828
client_streams = (server_to_client_receive, client_to_server_send)
2929
server_streams = (client_to_server_receive, server_to_client_send)
3030

31-
async with (
32-
server_to_client_receive,
33-
client_to_server_send,
34-
client_to_server_receive,
35-
server_to_client_send,
36-
):
31+
async with server_to_client_receive, client_to_server_send, client_to_server_receive, server_to_client_send:
3732
yield client_streams, server_streams

0 commit comments

Comments
 (0)