fix: use common parsing/context module between

NorthAuthenticationMiddleware and FastMCPNorthMiddleware

Author: speterkins

README.md

@@ -49,6 +49,8 @@ def echo(_: dict) -> dict:
 
 The middleware reads the `X-North-ID-Token` header (if present) and parses Base64-encoded JSON from `X-North-Connector-Tokens`. It never returns a 401—it simply exposes these values through a context variable and `request.state.north_context` for downstream handlers.
 
+When you use `NorthMCPServer`, the authentication stack now populates the same request context automatically, so utilities built against `get_north_request_context()` can be shared across FastMCP apps and fully authenticated servers.
+
 ## Examples
 
 This repository contains example servers that you can use as a quickstart. You can find them in the [examples directory](https://github.com/cohere-ai/north-mcp-python-sdk/tree/main/examples).

src/north_mcp_python_sdk/auth.py

@@ -18,6 +18,16 @@
 from starlette.responses import JSONResponse
 from starlette.types import ASGIApp, Receive, Scope, Send
 
+from .north_context import (
+    DEFAULT_CONNECTOR_TOKENS_HEADER,
+    DEFAULT_SERVER_SECRET_HEADER,
+    DEFAULT_USER_ID_TOKEN_HEADER,
+    NORTH_CONTEXT_SCOPE_KEY,
+    NorthRequestContext,
+    decode_connector_tokens,
+    reset_north_request_context,
+    set_north_request_context,
+)
 
 class AuthHeaderTokens(BaseModel):
     server_secret: str | None
@@ -30,9 +40,15 @@ def __init__(
         self,
         connector_access_tokens: dict[str, str],
         email: str | None = None,
+        user_id_token: str | None = None,
     ):
         self.connector_access_tokens = connector_access_tokens
         self.email = email
+        self.user_id_token = user_id_token
+        self.north_context = NorthRequestContext(
+            user_id_token=user_id_token,
+            connector_tokens=connector_access_tokens,
+        )
 
 
 class NorthAuthenticationMiddleware(AuthenticationMiddleware):
@@ -153,16 +169,32 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
             return await self.app(scope, receive, send)
 
         user = scope.get("user")
+        existing_context = scope.get(NORTH_CONTEXT_SCOPE_KEY)
+
+        def store_context(context: NorthRequestContext) -> None:
+            scope[NORTH_CONTEXT_SCOPE_KEY] = context
+            state = scope.get("state")
+            if state is None:
+                scope["state"] = {"north_context": context}
+            elif isinstance(state, dict):
+                state["north_context"] = context
+            else:
+                setattr(state, "north_context", context)
 
         # For custom routes that don't require auth, user will be None
         if user is None:
             self.logger.debug(
                 "Custom route accessed without authentication (operational endpoint)"
             )
+            context = existing_context or NorthRequestContext()
+            store_context(context)
+
+            context_token = set_north_request_context(context)
             token = auth_context_var.set(None)
             try:
                 await self.app(scope, receive, send)
             finally:
+                reset_north_request_context(context_token)
                 auth_context_var.reset(token)
             return
 
@@ -179,10 +211,15 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
             list(user.connector_access_tokens.keys()),
         )
 
+        context = existing_context or user.north_context
+        store_context(context)
+
+        context_token = set_north_request_context(context)
         token = auth_context_var.set(user)
         try:
             await self.app(scope, receive, send)
         finally:
+            reset_north_request_context(context_token)
             auth_context_var.reset(token)
 
 
@@ -211,26 +248,12 @@ def _has_x_north_headers(self, conn: HTTPConnection) -> bool:
         return any(
             conn.headers.get(header)
             for header in [
-                "X-North-ID-Token",
-                "X-North-Connector-Tokens",
-                "X-North-Server-Secret",
+                DEFAULT_USER_ID_TOKEN_HEADER,
+                DEFAULT_CONNECTOR_TOKENS_HEADER,
+                DEFAULT_SERVER_SECRET_HEADER,
             ]
         )
 
-    def _parse_connector_tokens(self, header_value: str) -> dict[str, str]:
-        """Parse Base64 URL-safe encoded JSON connector tokens."""
-        try:
-            # Add padding if needed for Base64 decoding
-            padded = header_value + "=" * (4 - len(header_value) % 4)
-            decoded_json = base64.urlsafe_b64decode(padded).decode()
-            tokens = json.loads(decoded_json)
-            if not isinstance(tokens, dict):
-                raise ValueError("Connector tokens must be a JSON object")
-            return tokens
-        except Exception as e:
-            self.logger.debug("Failed to parse connector tokens: %s", e)
-            raise AuthenticationError("invalid connector tokens format")
-
     def _validate_server_secret(self, provided_secret: str | None) -> None:
         """Validate server secret matches expected value."""
         if self._server_secret and self._server_secret != provided_secret:
@@ -271,11 +294,16 @@ def _process_user_id_token(self, user_id_token: str | None) -> str | None:
             raise AuthenticationError("invalid user id token")
 
     def _create_authenticated_user(
-        self, email: str | None, connector_access_tokens: dict[str, str]
+        self,
+        email: str | None,
+        connector_access_tokens: dict[str, str],
+        user_id_token: str | None,
     ) -> tuple[AuthCredentials, AuthenticatedNorthUser]:
         """Create authenticated user from validated tokens."""
         return AuthCredentials(), AuthenticatedNorthUser(
-            connector_access_tokens=connector_access_tokens, email=email
+            connector_access_tokens=connector_access_tokens,
+            email=email,
+            user_id_token=user_id_token,
         )
 
     async def _authenticate_x_north_headers(
@@ -285,16 +313,25 @@ async def _authenticate_x_north_headers(
         self.logger.debug("Using X-North headers for authentication")
 
         # Extract headers
-        user_id_token = conn.headers.get("X-North-ID-Token")
-        connector_tokens_header = conn.headers.get("X-North-Connector-Tokens")
-        server_secret = conn.headers.get("X-North-Server-Secret")
+        user_id_token = conn.headers.get(DEFAULT_USER_ID_TOKEN_HEADER)
+        connector_tokens_header = conn.headers.get(
+            DEFAULT_CONNECTOR_TOKENS_HEADER
+        )
+        server_secret = conn.headers.get(DEFAULT_SERVER_SECRET_HEADER)
 
         # Parse connector tokens (Base64 URL-safe encoded JSON)
         connector_access_tokens = {}
         if connector_tokens_header:
-            connector_access_tokens = self._parse_connector_tokens(
-                connector_tokens_header
-            )
+            try:
+                connector_access_tokens = decode_connector_tokens(
+                    connector_tokens_header,
+                    logger=self.logger,
+                    raise_on_error=True,
+                )
+            except ValueError as exc:
+                raise AuthenticationError(
+                    "invalid connector tokens format"
+                ) from exc
 
         self.logger.debug(
             "X-North headers parsed. Has server_secret: %s, Has user_id_token: %s, Connector count: %d",
@@ -309,8 +346,16 @@ async def _authenticate_x_north_headers(
         self._validate_server_secret(server_secret)
         email = self._process_user_id_token(user_id_token)
 
+        context = NorthRequestContext(
+            user_id_token=user_id_token,
+            connector_tokens=connector_access_tokens,
+        )
+        conn.scope[NORTH_CONTEXT_SCOPE_KEY] = context
+
         self.logger.debug("X-North authentication successful")
-        return self._create_authenticated_user(email, connector_access_tokens)
+        return self._create_authenticated_user(
+            email, connector_access_tokens, user_id_token
+        )
 
     async def _authenticate_legacy_bearer(
         self, conn: HTTPConnection
@@ -358,9 +403,15 @@ async def _authenticate_legacy_bearer(
         self._validate_server_secret(tokens.server_secret)
         email = self._process_user_id_token(tokens.user_id_token)
 
+        context = NorthRequestContext(
+            user_id_token=tokens.user_id_token,
+            connector_tokens=tokens.connector_access_tokens,
+        )
+        conn.scope[NORTH_CONTEXT_SCOPE_KEY] = context
+
         self.logger.debug("Legacy authentication successful")
         return self._create_authenticated_user(
-            email, tokens.connector_access_tokens
+            email, tokens.connector_access_tokens, tokens.user_id_token
         )
 
     async def authenticate(

src/north_mcp_python_sdk/middleware.py

@@ -1,44 +1,21 @@
-import base64
-import binascii
-import contextvars
-import json
 import logging
-from dataclasses import dataclass, field
-from typing import Dict, Optional
 
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.requests import Request
 from starlette.types import ASGIApp
 
-
-_DEFAULT_USER_ID_TOKEN_HEADER = "X-North-ID-Token"
-_DEFAULT_CONNECTOR_TOKENS_HEADER = "X-North-Connector-Tokens"
-
-
-@dataclass(frozen=True)
-class NorthRequestContext:
-    """Holds North-specific request context extracted from headers."""
-
-    user_id_token: Optional[str] = None
-    connector_tokens: Dict[str, str] = field(default_factory=dict)
-
-
-north_request_context_var = contextvars.ContextVar[NorthRequestContext](
-    "north_request_context",
-    default=NorthRequestContext(),
+from .north_context import (
+    DEFAULT_CONNECTOR_TOKENS_HEADER,
+    DEFAULT_USER_ID_TOKEN_HEADER,
+    NorthRequestContext,
+    decode_connector_tokens,
+    get_north_request_context,
+    north_request_context_var,
+    reset_north_request_context,
+    set_north_request_context,
 )
 
 
-def get_north_request_context() -> NorthRequestContext:
-    """
-    Retrieve the North request context for the current request.
-
-    Returns:
-        NorthRequestContext: The context extracted by FastMCPNorthMiddleware.
-    """
-    return north_request_context_var.get()
-
-
 class FastMCPNorthMiddleware(BaseHTTPMiddleware):
     """
     Lightweight middleware that extracts North request metadata from headers.
@@ -52,8 +29,8 @@ def __init__(
         self,
         app: ASGIApp,
         *,
-        user_id_token_header: str = _DEFAULT_USER_ID_TOKEN_HEADER,
-        connector_tokens_header: str = _DEFAULT_CONNECTOR_TOKENS_HEADER,
+        user_id_token_header: str = DEFAULT_USER_ID_TOKEN_HEADER,
+        connector_tokens_header: str = DEFAULT_CONNECTOR_TOKENS_HEADER,
         debug: bool = False,
     ) -> None:
         super().__init__(app)
@@ -64,43 +41,13 @@ def __init__(
             self._logger.setLevel(logging.DEBUG)
         self._debug = debug
 
-    def _parse_connector_tokens(self, raw_header: str) -> Dict[str, str]:
-        """
-        Parse the connector tokens header, expected to be Base64-encoded JSON.
-
-        Returns an empty dict when the header cannot be decoded or does not
-        resolve to a JSON object of string keys and values.
-        """
-        if not raw_header:
-            return {}
-
-        padding = (-len(raw_header)) % 4
-        padded_value = raw_header + ("=" * padding)
-
-        try:
-            decoded_bytes = base64.urlsafe_b64decode(padded_value)
-            decoded_json = decoded_bytes.decode()
-            parsed = json.loads(decoded_json)
-            if isinstance(parsed, dict):
-                return {
-                    str(key): str(value)
-                    for key, value in parsed.items()
-                    if isinstance(key, str) and isinstance(value, str)
-                }
-        except (ValueError, json.JSONDecodeError, binascii.Error) as exc:
-            self._logger.debug(
-                "Failed to decode connector tokens header: %s", exc
-            )
-
-        return {}
-
     async def dispatch(self, request: Request, call_next):
         user_id_token = request.headers.get(self._user_id_token_header)
         connector_tokens_header = request.headers.get(
             self._connector_tokens_header
         )
-        connector_tokens = self._parse_connector_tokens(
-            connector_tokens_header or ""
+        connector_tokens = decode_connector_tokens(
+            connector_tokens_header or "", logger=self._logger
         )
 
         if self._debug:
@@ -116,12 +63,12 @@ async def dispatch(self, request: Request, call_next):
         )
 
         request.state.north_context = context
-        token = north_request_context_var.set(context)
+        token = set_north_request_context(context)
 
         try:
             response = await call_next(request)
         finally:
-            north_request_context_var.reset(token)
+            reset_north_request_context(token)
             # Request.state lives for the lifetime of the request; no cleanup needed.
 
         return response