Merge pull request #16 from sacha-development-stuff/codex/resolve-merge-conflicts-and-check-functionality

Fix merge conflicts and restore auth features

Author: SoldierSacha

src/mcp/client/auth.py

@@ -486,6 +486,8 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                     # Retry with new tokens
                     self._add_auth_header(request)
                     yield request
+
+
 class ClientCredentialsProvider(httpx.Auth):
     """HTTPX auth using the OAuth2 client credentials grant."""
 
@@ -508,19 +510,48 @@ def __init__(
 
         self._token_lock = anyio.Lock()
 
+    def _get_authorization_base_url(self, server_url: str) -> str:
+        """Return base authorization server URL without path."""
+        parsed = urlparse(server_url)
+        return f"{parsed.scheme}://{parsed.netloc}"
+
+    async def _discover_oauth_metadata(self, server_url: str) -> OAuthMetadata | None:
+        """Discover OAuth server metadata for client credentials."""
+        auth_base_url = self._get_authorization_base_url(server_url)
+        url = urljoin(auth_base_url, "/.well-known/oauth-authorization-server")
+        headers = {"MCP-Protocol-Version": LATEST_PROTOCOL_VERSION}
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=headers)
+                if response.status_code == 404:
+                    return None
+                response.raise_for_status()
+                return OAuthMetadata.model_validate(response.json())
+            except Exception:
+                try:
+                    response = await client.get(url)
+                    if response.status_code == 404:
+                        return None
+                    response.raise_for_status()
+                    return OAuthMetadata.model_validate(response.json())
+                except Exception:
+                    logger.exception("Failed to discover OAuth metadata")
+                    return None
+
     async def _register_oauth_client(
         self,
         server_url: str,
         client_metadata: OAuthClientMetadata,
         metadata: OAuthMetadata | None = None,
     ) -> OAuthClientInformationFull:
         if not metadata:
-            metadata = await _discover_oauth_metadata(server_url)
+            metadata = await self._discover_oauth_metadata(server_url)
 
         if metadata and metadata.registration_endpoint:
             registration_url = str(metadata.registration_endpoint)
         else:
-            auth_base_url = _get_authorization_base_url(server_url)
+            auth_base_url = self._get_authorization_base_url(server_url)
             registration_url = urljoin(auth_base_url, "/register")
 
         if client_metadata.scope is None and metadata and metadata.scopes_supported is not None:
@@ -582,14 +613,14 @@ async def _get_or_register_client(self) -> OAuthClientInformationFull:
 
     async def _request_token(self) -> None:
         if not self._metadata:
-            self._metadata = await _discover_oauth_metadata(self.server_url)
+            self._metadata = await self._discover_oauth_metadata(self.server_url)
 
         client_info = await self._get_or_register_client()
 
         if self._metadata and self._metadata.token_endpoint:
             token_url = str(self._metadata.token_endpoint)
         else:
-            auth_base_url = _get_authorization_base_url(self.server_url)
+            auth_base_url = self._get_authorization_base_url(self.server_url)
             token_url = urljoin(auth_base_url, "/token")
 
         token_data = {
@@ -671,14 +702,14 @@ def __init__(
 
     async def _request_token(self) -> None:
         if not self._metadata:
-            self._metadata = await _discover_oauth_metadata(self.server_url)
+            self._metadata = await self._discover_oauth_metadata(self.server_url)
 
         client_info = await self._get_or_register_client()
 
         if self._metadata and self._metadata.token_endpoint:
             token_url = str(self._metadata.token_endpoint)
         else:
-            auth_base_url = _get_authorization_base_url(self.server_url)
+            auth_base_url = self._get_authorization_base_url(self.server_url)
             token_url = urljoin(auth_base_url, "/token")
 
         subject_token = await self.subject_token_supplier()

tests/client/test_auth.py

@@ -2,18 +2,24 @@
 Tests for refactored OAuth client authentication implementation.
 """
 
-import time
 import asyncio
+import time
+from unittest.mock import AsyncMock, Mock, patch
 
 import httpx
 import pytest
 from pydantic import AnyHttpUrl, AnyUrl
-from unittest.mock import AsyncMock, Mock, patch
 
-from mcp.client.auth import OAuthClientProvider, PKCEParameters
+from mcp.client.auth import (
+    ClientCredentialsProvider,
+    OAuthClientProvider,
+    PKCEParameters,
+    TokenExchangeProvider,
+)
 from mcp.shared.auth import (
     OAuthClientInformationFull,
     OAuthClientMetadata,
+    OAuthMetadata,
     OAuthToken,
 )
 
@@ -81,6 +87,8 @@ async def callback_handler() -> tuple[str, str | None]:
         redirect_handler=redirect_handler,
         callback_handler=callback_handler,
     )
+
+
 @pytest.fixture
 def client_credentials_metadata():
     return OAuthClientMetadata(
@@ -92,6 +100,45 @@ def client_credentials_metadata():
         token_endpoint_auth_method="client_secret_post",
     )
 
+
+@pytest.fixture
+def oauth_metadata():
+    return OAuthMetadata(
+        issuer=AnyHttpUrl("https://auth.example.com"),
+        authorization_endpoint=AnyHttpUrl("https://auth.example.com/authorize"),
+        token_endpoint=AnyHttpUrl("https://auth.example.com/token"),
+        registration_endpoint=AnyHttpUrl("https://auth.example.com/register"),
+        scopes_supported=["read", "write", "admin"],
+        response_types_supported=["code"],
+        grant_types_supported=["authorization_code", "refresh_token", "client_credentials"],
+        code_challenge_methods_supported=["S256"],
+    )
+
+
+@pytest.fixture
+def oauth_client_info():
+    return OAuthClientInformationFull(
+        client_id="test_client_id",
+        client_secret="test_client_secret",
+        redirect_uris=[AnyUrl("http://localhost:3000/callback")],
+        client_name="Test Client",
+        grant_types=["authorization_code", "refresh_token"],
+        response_types=["code"],
+        scope="read write",
+    )
+
+
+@pytest.fixture
+def oauth_token():
+    return OAuthToken(
+        access_token="test_access_token",
+        token_type="bearer",
+        expires_in=3600,
+        refresh_token="test_refresh_token",
+        scope="read write",
+    )
+
+
 @pytest.fixture
 async def client_credentials_provider(client_credentials_metadata, mock_storage):
     return ClientCredentialsProvider(
@@ -100,6 +147,7 @@ async def client_credentials_provider(client_credentials_metadata, mock_storage)
         storage=mock_storage,
     )
 
+
 @pytest.fixture
 async def token_exchange_provider(client_credentials_metadata, mock_storage):
     return TokenExchangeProvider(
@@ -342,6 +390,8 @@ async def test_auth_flow_with_valid_tokens(self, oauth_provider, mock_storage, v
             await auth_flow.asend(response)
         except StopAsyncIteration:
             pass  # Expected
+
+
 class TestClientCredentialsProvider:
     @pytest.mark.anyio
     async def test_request_token_success(
@@ -417,4 +467,3 @@ async def test_request_token_success(
 
             mock_client.post.assert_called_once()
             assert token_exchange_provider._current_tokens.access_token == oauth_token.access_token
-