chore(auth): type-safe auth context push

Author: Epochex

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

@@ -1,6 +1,6 @@
 import contextvars
-
 from contextvars import Token
+from typing import Any
 
 from starlette.requests import Request
 from starlette.types import ASGIApp, Receive, Scope, Send
@@ -23,7 +23,7 @@ def get_access_token() -> AccessToken | None:
     return auth_user.access_token if auth_user else None
 
 
-def _push_auth_context_from_request(request: Request | None) -> Token[AuthenticatedUser | None] | None:
+def push_auth_context_from_request(request: Request | None) -> Token[AuthenticatedUser | None] | None:
     """Set auth context for the current task from an incoming request.
 
     This is primarily used by server transports where request handlers may run
@@ -32,10 +32,7 @@ def _push_auth_context_from_request(request: Request | None) -> Token[Authentica
     if request is None:
         return None
     # Avoid Request.user, which asserts AuthenticationMiddleware is installed.
-    user = None
-    scope = getattr(request, "scope", None)
-    if isinstance(scope, dict):
-        user = scope.get("user")
+    user: Any | None = request.scope.get("user")
     if user is None:
         try:
             user = getattr(request, "user", None)
@@ -46,7 +43,7 @@ def _push_auth_context_from_request(request: Request | None) -> Token[Authentica
     return None
 
 
-def _pop_auth_context(token: Token[AuthenticatedUser | None] | None) -> None:
+def pop_auth_context(token: Token[AuthenticatedUser | None] | None) -> None:
     if token is None:
         return
     auth_context_var.reset(token)

src/mcp/server/lowlevel/server.py

@@ -53,7 +53,11 @@ async def main():
 from typing_extensions import TypeVar
 
 from mcp import types
-from mcp.server.auth.middleware.auth_context import AuthContextMiddleware, _pop_auth_context, _push_auth_context_from_request
+from mcp.server.auth.middleware.auth_context import (
+    AuthContextMiddleware,
+    pop_auth_context,
+    push_auth_context_from_request,
+)
 from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend, RequireAuthMiddleware
 from mcp.server.auth.provider import OAuthAuthorizationServerProvider, TokenVerifier
 from mcp.server.auth.routes import build_resource_metadata_url, create_auth_routes, create_protected_resource_routes
@@ -497,11 +501,11 @@ async def _handle_request(
                         close_sse_stream=close_sse_stream_cb,
                         close_standalone_sse_stream=close_standalone_sse_stream_cb,
                     )
-                    auth_token = _push_auth_context_from_request(request_data)
+                    auth_token = push_auth_context_from_request(request_data)
                     try:
                         response = await handler(ctx, req.params)
                     finally:
-                        _pop_auth_context(auth_token)
+                        pop_auth_context(auth_token)
                 except MCPError as err:
                     response = err.error
                 except anyio.get_cancelled_exc_class():