refactor

Author: ihrpr

examples/clients/simple-auth-client/mcp_simple_auth_client/main.py

@@ -15,7 +15,7 @@
 from typing import Any
 from urllib.parse import parse_qs, urlparse
 
-from mcp.client.oauth_auth import OAuthAuth
+from mcp.client.auth import OAuthAuth
 from mcp.client.session import ClientSession
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.shared.auth import OAuthClientMetadata

src/mcp/client/oauth_auth.py renamed to src/mcp/client/auth.py

@@ -2,7 +2,8 @@
 Production-ready OAuth2 Authentication implementation for HTTPX using anyio.
 
 This module provides a complete OAuth 2.0 authentication implementation
-that handles authorization code flow with PKCE, automatic token refresh and proper error handling.
+that handles authorization code flow with PKCE,
+automatic token refresh and proper error handling.
 The callback server implementation should be handled by the calling code.
 """
 
@@ -12,13 +13,13 @@
 import string
 import time
 import webbrowser
-from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from urllib.parse import urljoin
 
 import anyio
 import httpx
 
+from mcp.client.token_storage import InMemoryTokenStorage, TokenStorage
 from mcp.shared.auth import (
     OAuthClientInformationFull,
     OAuthClientMetadata,
@@ -28,50 +29,6 @@
 from mcp.types import LATEST_PROTOCOL_VERSION
 
 
-class TokenStorage(ABC):
-    """Abstract base class for token storage implementations."""
-
-    @abstractmethod
-    async def get_tokens(self) -> OAuthToken | None:
-        """Get stored tokens."""
-        pass
-
-    @abstractmethod
-    async def set_tokens(self, tokens: OAuthToken) -> None:
-        """Store tokens."""
-        pass
-
-    @abstractmethod
-    async def get_client_info(self) -> OAuthClientInformationFull | None:
-        """Get stored client information."""
-        pass
-
-    @abstractmethod
-    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
-        """Store client information."""
-        pass
-
-
-class InMemoryTokenStorage(TokenStorage):
-    """Simple in-memory token storage implementation."""
-
-    def __init__(self):
-        self._tokens: OAuthToken | None = None
-        self._client_info: OAuthClientInformationFull | None = None
-
-    async def get_tokens(self) -> OAuthToken | None:
-        return self._tokens
-
-    async def set_tokens(self, tokens: OAuthToken) -> None:
-        self._tokens = tokens
-
-    async def get_client_info(self) -> OAuthClientInformationFull | None:
-        return self._client_info
-
-    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
-        self._client_info = client_info
-
-
 async def discover_oauth_metadata(server_url: str) -> OAuthMetadata | None:
     """
     Discovers OAuth metadata from the server's well-known endpoint.

src/mcp/client/token_storage.py

@@ -0,0 +1,50 @@
+"""
+Token storage implementations for OAuth authentication.
+
+This module provides protocol and concrete implementations for storing
+OAuth tokens and client information.
+"""
+
+from typing import Protocol
+
+from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
+
+
+class TokenStorage(Protocol):
+    """Protocol for token storage implementations."""
+
+    async def get_tokens(self) -> OAuthToken | None:
+        """Get stored tokens."""
+        ...
+
+    async def set_tokens(self, tokens: OAuthToken) -> None:
+        """Store tokens."""
+        ...
+
+    async def get_client_info(self) -> OAuthClientInformationFull | None:
+        """Get stored client information."""
+        ...
+
+    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        """Store client information."""
+        ...
+
+
+class InMemoryTokenStorage(TokenStorage):
+    """Simple in-memory token storage implementation."""
+
+    def __init__(self):
+        self._tokens: OAuthToken | None = None
+        self._client_info: OAuthClientInformationFull | None = None
+
+    async def get_tokens(self) -> OAuthToken | None:
+        return self._tokens
+
+    async def set_tokens(self, tokens: OAuthToken) -> None:
+        self._tokens = tokens
+
+    async def get_client_info(self) -> OAuthClientInformationFull | None:
+        return self._client_info
+
+    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        self._client_info = client_info