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/runner.py

@@ -26,7 +26,7 @@
 from pydantic import BaseModel, ValidationError
 from typing_extensions import TypeVar
 
-from mcp.server.auth.middleware.auth_context import _pop_auth_context, _push_auth_context_from_request
+from mcp.server.auth.middleware.auth_context import pop_auth_context, push_auth_context_from_request
 from mcp.server.connection import Connection
 from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, ServerRequestContext
 from mcp.server.models import InitializationOptions
@@ -260,11 +260,11 @@ async def _inner() -> HandlerResult:
             return result
 
         call = self._compose_server_middleware(ctx, method, params, _inner)
-        auth_token = _push_auth_context_from_request(ctx.request)
+        auth_token = push_auth_context_from_request(ctx.request)
         try:
             result = _dump_result(await call())
         finally:
-            _pop_auth_context(auth_token)
+            pop_auth_context(auth_token)
         if method == "initialize":
             # Commit only on chain success, so a middleware veto leaves no state.
             # Race-free: the read loop is parked until this call returns.