mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 13:34:17 +02:00
Allow OAuth instance to use the same httpx factory as the Transport (#2324)
* Allow OAuth instance to use the same httpx factory as the Transport * Fix test * Update SSL verification mode assertion in tests * This is actually not needed * Creating a Client instance is not needed for this test * Fix test * Apply httpx_client_factory fix to SSETransport Extends the OAuth httpx_client_factory changes to SSETransport. SSETransport had the same issues as StreamableHttpTransport where it wasn't passing the custom httpx client factory to OAuth, causing certificate verification settings to be ignored during OAuth flows. Changes: - Set httpx_client_factory before calling _set_auth() - Pass httpx_client_factory to OAuth constructor - Add test for SSETransport OAuth client factory propagation --------- Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
This commit is contained in:
parent
a7563a39b0
commit
a6ddde27df
3 changed files with 45 additions and 5 deletions
|
|
@ -12,6 +12,7 @@ from key_value.aio.adapters.pydantic import PydanticAdapter
|
|||
from key_value.aio.protocols import AsyncKeyValue
|
||||
from key_value.aio.stores.memory import MemoryStore
|
||||
from mcp.client.auth import OAuthClientProvider, TokenStorage
|
||||
from mcp.shared._httpx_utils import McpHttpClientFactory
|
||||
from mcp.shared.auth import (
|
||||
OAuthClientInformationFull,
|
||||
OAuthClientMetadata,
|
||||
|
|
@ -147,6 +148,7 @@ class OAuth(OAuthClientProvider):
|
|||
token_storage: AsyncKeyValue | None = None,
|
||||
additional_client_metadata: dict[str, Any] | None = None,
|
||||
callback_port: int | None = None,
|
||||
httpx_client_factory: McpHttpClientFactory | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize OAuth client provider for an MCP server.
|
||||
|
|
@ -164,6 +166,7 @@ class OAuth(OAuthClientProvider):
|
|||
server_base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
|
||||
|
||||
# Setup OAuth client
|
||||
self.httpx_client_factory = httpx_client_factory or httpx.AsyncClient
|
||||
self.redirect_port = callback_port or find_available_port()
|
||||
redirect_uri = f"http://localhost:{self.redirect_port}/callback"
|
||||
|
||||
|
|
@ -226,7 +229,7 @@ class OAuth(OAuthClientProvider):
|
|||
async def redirect_handler(self, authorization_url: str) -> None:
|
||||
"""Open browser for authorization, with pre-flight check for invalid client."""
|
||||
# Pre-flight check to detect invalid client_id before opening browser
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self.httpx_client_factory() as client:
|
||||
response = await client.get(authorization_url, follow_redirects=False)
|
||||
|
||||
# Check for client not found error (400 typically means bad client_id)
|
||||
|
|
|
|||
|
|
@ -177,8 +177,8 @@ class SSETransport(ClientTransport):
|
|||
|
||||
self.url = url
|
||||
self.headers = headers or {}
|
||||
self._set_auth(auth)
|
||||
self.httpx_client_factory = httpx_client_factory
|
||||
self._set_auth(auth)
|
||||
|
||||
if isinstance(sse_read_timeout, int | float):
|
||||
sse_read_timeout = datetime.timedelta(seconds=float(sse_read_timeout))
|
||||
|
|
@ -186,7 +186,7 @@ class SSETransport(ClientTransport):
|
|||
|
||||
def _set_auth(self, auth: httpx.Auth | Literal["oauth"] | str | None):
|
||||
if auth == "oauth":
|
||||
auth = OAuth(self.url)
|
||||
auth = OAuth(self.url, httpx_client_factory=self.httpx_client_factory)
|
||||
elif isinstance(auth, str):
|
||||
auth = BearerAuth(auth)
|
||||
self.auth = auth
|
||||
|
|
@ -247,8 +247,8 @@ class StreamableHttpTransport(ClientTransport):
|
|||
|
||||
self.url = url
|
||||
self.headers = headers or {}
|
||||
self._set_auth(auth)
|
||||
self.httpx_client_factory = httpx_client_factory
|
||||
self._set_auth(auth)
|
||||
|
||||
if isinstance(sse_read_timeout, int | float):
|
||||
sse_read_timeout = datetime.timedelta(seconds=float(sse_read_timeout))
|
||||
|
|
@ -256,7 +256,7 @@ class StreamableHttpTransport(ClientTransport):
|
|||
|
||||
def _set_auth(self, auth: httpx.Auth | Literal["oauth"] | str | None):
|
||||
if auth == "oauth":
|
||||
auth = OAuth(self.url)
|
||||
auth = OAuth(self.url, httpx_client_factory=self.httpx_client_factory)
|
||||
elif isinstance(auth, str):
|
||||
auth = BearerAuth(auth)
|
||||
self.auth = auth
|
||||
|
|
|
|||
37
tests/client/transports/test_transports.py
Normal file
37
tests/client/transports/test_transports.py
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
from ssl import VerifyMode
|
||||
|
||||
import httpx
|
||||
|
||||
from fastmcp.client.transports import SSETransport, StreamableHttpTransport
|
||||
|
||||
|
||||
async def test_oauth_uses_same_client_as_transport_streamable_http():
|
||||
transport = StreamableHttpTransport(
|
||||
"https://some.fake.url/",
|
||||
httpx_client_factory=lambda *args, **kwargs: httpx.AsyncClient(
|
||||
verify=False, *args, **kwargs
|
||||
),
|
||||
auth="oauth",
|
||||
)
|
||||
|
||||
async with transport.auth.httpx_client_factory() as httpx_client: # type: ignore[attr-defined]
|
||||
assert (
|
||||
httpx_client._transport._pool._ssl_context.verify_mode
|
||||
== VerifyMode.CERT_NONE
|
||||
)
|
||||
|
||||
|
||||
async def test_oauth_uses_same_client_as_transport_sse():
|
||||
transport = SSETransport(
|
||||
"https://some.fake.url/",
|
||||
httpx_client_factory=lambda *args, **kwargs: httpx.AsyncClient(
|
||||
verify=False, *args, **kwargs
|
||||
),
|
||||
auth="oauth",
|
||||
)
|
||||
|
||||
async with transport.auth.httpx_client_factory() as httpx_client: # type: ignore[attr-defined]
|
||||
assert (
|
||||
httpx_client._transport._pool._ssl_context.verify_mode
|
||||
== VerifyMode.CERT_NONE
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue