Typecheck

Author: praboud-ant

src/mcp/client/sse.py

@@ -1,6 +1,6 @@
 import logging
 from contextlib import asynccontextmanager
-from typing import Any, Union
+from typing import Any
 from urllib.parse import urljoin, urlparse
 
 import anyio
@@ -26,7 +26,7 @@ async def sse_client(
     headers: dict[str, Any] | None = None,
     timeout: float = 5,
     sse_read_timeout: float = 60 * 5,
-    auth: Union[AuthSession, OAuthClient, None] = None,
+    auth: AuthSession | OAuthClient | None = None,
 ):
     """
     Client transport for SSE.

src/mcp/server/auth/handlers/authorize.py

@@ -1,6 +1,6 @@
 import logging
 from dataclasses import dataclass
-from typing import Literal
+from typing import Any, Literal
 from urllib.parse import urlencode, urlparse, urlunparse
 
 from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, RootModel, ValidationError
@@ -70,13 +70,13 @@ def best_effort_extract_string(
     return None
 
 
-class AnyHttpUrlModel(RootModel):
+class AnyHttpUrlModel(RootModel[AnyHttpUrl]):
     root: AnyHttpUrl
 
 
 @dataclass
 class AuthorizationHandler:
-    provider: OAuthServerProvider
+    provider: OAuthServerProvider[Any, Any, Any]
 
     async def handle(self, request: Request) -> Response:
         # implements authorization requests for grant_type=code;

src/mcp/server/auth/handlers/register.py

@@ -1,6 +1,7 @@
 import secrets
 import time
 from dataclasses import dataclass
+from typing import Any
 from uuid import uuid4
 
 from pydantic import BaseModel, RootModel, ValidationError
@@ -18,7 +19,7 @@
 from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
 
 
-class RegistrationRequest(RootModel):
+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
@@ -31,7 +32,7 @@ class RegistrationErrorResponse(BaseModel):
 
 @dataclass
 class RegistrationHandler:
-    provider: OAuthServerProvider
+    provider: OAuthServerProvider[Any, Any, Any]
     options: ClientRegistrationOptions
 
     async def handle(self, request: Request) -> Response:

src/mcp/server/auth/handlers/revoke.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from functools import partial
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import BaseModel, ValidationError
 from starlette.requests import Request
@@ -35,7 +35,7 @@ class RevocationErrorResponse(BaseModel):
 
 @dataclass
 class RevocationHandler:
-    provider: OAuthServerProvider
+    provider: OAuthServerProvider[Any, Any, Any]
     client_authenticator: ClientAuthenticator
 
     async def handle(self, request: Request) -> Response:

src/mcp/server/auth/handlers/token.py

@@ -2,7 +2,7 @@
 import hashlib
 import time
 from dataclasses import dataclass
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 from pydantic import AnyHttpUrl, BaseModel, Field, RootModel, ValidationError
 from starlette.requests import Request
@@ -44,7 +44,14 @@ class RefreshTokenRequest(BaseModel):
     client_secret: str | None = None
 
 
-class TokenRequest(RootModel):
+class TokenRequest(
+    RootModel[
+        Annotated[
+            AuthorizationCodeRequest | RefreshTokenRequest,
+            Field(discriminator="grant_type"),
+        ]
+    ]
+):
     root: Annotated[
         AuthorizationCodeRequest | RefreshTokenRequest,
         Field(discriminator="grant_type"),
@@ -61,7 +68,7 @@ class TokenErrorResponse(BaseModel):
     error_uri: AnyHttpUrl | None = None
 
 
-class TokenSuccessResponse(RootModel):
+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
@@ -70,7 +77,7 @@ class TokenSuccessResponse(RootModel):
 
 @dataclass
 class TokenHandler:
-    provider: OAuthServerProvider
+    provider: OAuthServerProvider[Any, Any, Any]
     client_authenticator: ClientAuthenticator
 
     def response(self, obj: TokenSuccessResponse | TokenErrorResponse | ErrorResponse):

src/mcp/server/auth/middleware/bearer_auth.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any, Callable
+from typing import Any
 
 from starlette.authentication import (
     AuthCredentials,
@@ -8,7 +8,7 @@
 )
 from starlette.exceptions import HTTPException
 from starlette.requests import HTTPConnection
-from starlette.types import Scope
+from starlette.types import Receive, Scope, Send
 
 from mcp.server.auth.provider import AccessToken, OAuthServerProvider
 
@@ -29,7 +29,7 @@ class BearerAuthBackend(AuthenticationBackend):
 
     def __init__(
         self,
-        provider: OAuthServerProvider,
+        provider: OAuthServerProvider[Any, Any, Any],
     ):
         self.provider = provider
 
@@ -72,7 +72,7 @@ def __init__(self, app: Any, required_scopes: list[str]):
         self.app = app
         self.required_scopes = required_scopes
 
-    async def __call__(self, scope: Scope, receive: Callable, send: Callable) -> None:
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         auth_credentials = scope.get("auth")
 
         for required_scope in self.required_scopes:

src/mcp/server/auth/middleware/client_auth.py

@@ -1,4 +1,5 @@
 import time
+from typing import Any
 
 from mcp.server.auth.provider import OAuthServerProvider
 from mcp.shared.auth import OAuthClientInformationFull
@@ -20,7 +21,7 @@ class ClientAuthenticator:
     logic is skipped.
     """
 
-    def __init__(self, provider: OAuthServerProvider):
+    def __init__(self, provider: OAuthServerProvider[Any, Any, Any]):
         """
         Initialize the dependency.
 

src/mcp/server/auth/routes.py

@@ -1,4 +1,5 @@
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 from pydantic import AnyHttpUrl
 from starlette.routing import Route

src/mcp/server/fastmcp/server.py

@@ -5,13 +5,13 @@
 import inspect
 import json
 import re
-from collections.abc import AsyncIterator, Callable, Iterable, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
 from contextlib import (
     AbstractAsyncContextManager,
     asynccontextmanager,
 )
 from itertools import chain
-from typing import Any, Awaitable, Generic, Literal
+from typing import Any, Generic, Literal
 
 import anyio
 import pydantic_core
@@ -22,10 +22,10 @@
 from sse_starlette import EventSourceResponse
 from starlette.applications import Starlette
 from starlette.authentication import requires
+from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.requests import Request
 from starlette.responses import Response
-from starlette.middleware import Middleware
 from starlette.routing import Mount, Route
 
 from mcp.server.auth.middleware.auth_context import AuthContextMiddleware
@@ -491,7 +491,6 @@ def custom_route(
         name: str | None = None,
         include_in_schema: bool = True,
     ):
-
         def decorator(
             func: Callable[[Request], Awaitable[Response]],
         ) -> Callable[[Request], Awaitable[Response]]:
@@ -541,7 +540,7 @@ async def handle_sse(request: Request) -> EventSourceResponse:
             async with sse.connect_sse(
                 request.scope,
                 request.receive,
-                request._send # type: ignore[reportPrivateUsage]
+                request._send,  # type: ignore[reportPrivateUsage]
             ) as streams:
                 await self._mcp_server.run(
                     streams[0],
@@ -586,7 +585,9 @@ async def handle_sse(request: Request) -> EventSourceResponse:
 
         routes.append(
             Route(
-                self.settings.sse_path, endpoint=requires(required_scopes)(handle_sse), methods=["GET"]
+                self.settings.sse_path,
+                endpoint=requires(required_scopes)(handle_sse),
+                methods=["GET"],
             )
         )
         routes.append(
@@ -754,9 +755,9 @@ async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContent
         Returns:
             The resource content as either text or bytes
         """
-        assert self._fastmcp is not None, (
-            "Context is not available outside of a request"
-        )
+        assert (
+            self._fastmcp is not None
+        ), "Context is not available outside of a request"
         return await self._fastmcp.read_resource(uri)
 
     async def log(

src/mcp/server/streaming_asgi_transport.py

@@ -9,14 +9,15 @@
 """
 
 import typing
-from typing import Any, Dict, Tuple
+from typing import Any, cast
 
 import anyio
 import anyio.abc
 import anyio.streams.memory
 from httpx._models import Request, Response
 from httpx._transports.base import AsyncBaseTransport
 from httpx._types import AsyncByteStream
+from starlette.types import ASGIApp, Receive, Scope, Send
 
 
 class StreamingASGITransport(AsyncBaseTransport):
@@ -42,11 +43,11 @@ class StreamingASGITransport(AsyncBaseTransport):
 
     def __init__(
         self,
-        app: typing.Callable,
+        app: ASGIApp,
         task_group: anyio.abc.TaskGroup,
         raise_app_exceptions: bool = True,
         root_path: str = "",
-        client: Tuple[str, int] = ("127.0.0.1", 123),
+        client: tuple[str, int] = ("127.0.0.1", 123),
     ) -> None:
         self.app = app
         self.raise_app_exceptions = raise_app_exceptions
@@ -88,13 +89,15 @@ async def handle_async_request(
         initial_response_ready = anyio.Event()
 
         # Synchronization for streaming response
-        asgi_send_channel, asgi_receive_channel = anyio.create_memory_object_stream(100)
+        asgi_send_channel, asgi_receive_channel = anyio.create_memory_object_stream[
+            dict[str, Any]
+        ](100)
         content_send_channel, content_receive_channel = (
             anyio.create_memory_object_stream[bytes](100)
         )
 
         # ASGI callables.
-        async def receive() -> Dict[str, Any]:
+        async def receive() -> dict[str, Any]:
             nonlocal request_complete
 
             if request_complete:
@@ -108,15 +111,18 @@ async def receive() -> Dict[str, Any]:
                 return {"type": "http.request", "body": b"", "more_body": False}
             return {"type": "http.request", "body": body, "more_body": True}
 
-        async def send(message: Dict[str, Any]) -> None:
+        async def send(message: dict[str, Any]) -> None:
             nonlocal status_code, response_headers, response_started
 
             await asgi_send_channel.send(message)
 
         # Start the ASGI application in a separate task
         async def run_app() -> None:
             try:
-                await self.app(scope, receive, send)
+                # Cast the receive and send functions to the ASGI types
+                await self.app(
+                    cast(Scope, scope), cast(Receive, receive), cast(Send, send)
+                )
             except Exception:
                 if self.raise_app_exceptions:
                     raise