diff --git a/packages/mcp/pyproject.toml b/packages/mcp/pyproject.toml index 7dc621ee..efd481ee 100644 --- a/packages/mcp/pyproject.toml +++ b/packages/mcp/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "mcp>=2.0.0,<3.0", "pydantic>=2.11.7", "httpx>=0.27.2", + "httpx2>=2.5.0", "nanoid>=2.0.0", "aiohttp>=3.11.11", "aiosqlite>=0.20.0", diff --git a/packages/mcp/src/keycardai/mcp/client/auth/oauth/discovery.py b/packages/mcp/src/keycardai/mcp/client/auth/oauth/discovery.py index ef704c2a..0735a469 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/oauth/discovery.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/oauth/discovery.py @@ -4,6 +4,7 @@ from typing import Any from httpx import AsyncClient, Response +from httpx2 import Response as Httpx2Response from ...logging_config import get_logger from ..storage_facades import OAuthStorage @@ -38,7 +39,10 @@ def __init__( self.storage = storage self.client_factory = client_factory or default_client_factory - async def discover_resource(self, challenge_response: Response) -> dict[str, Any]: + async def discover_resource( + self, + challenge_response: Response | Httpx2Response, + ) -> dict[str, Any]: """ Discover protected resource metadata from 401 challenge. @@ -149,4 +153,3 @@ async def discover_auth_server( continue raise ValueError("Failed to discover any authorization server") - diff --git a/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py b/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py index c12e8204..b80e65a6 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Protocol from httpx import AsyncClient, Response +from httpx2 import Response as Httpx2Response from ...logging_config import get_logger from ...storage import NamespacedStorage @@ -225,7 +226,7 @@ async def get_auth_metadata(self) -> dict[str, Any]: async def handle_challenge( self, - challenge: Response, + challenge: Response | Httpx2Response, resource_url: str ) -> bool: """ @@ -244,7 +245,10 @@ async def handle_challenge( Returns: True if challenge was handled successfully, False otherwise """ - if not isinstance(challenge, Response) or challenge.status_code != 401: + if ( + not isinstance(challenge, (Response, Httpx2Response)) + or challenge.status_code != 401 + ): return False try: @@ -374,4 +378,3 @@ def create_auth_strategy( ) else: raise ValueError(f"Unknown auth type: {auth_type}") - diff --git a/packages/mcp/src/keycardai/mcp/client/auth/transports.py b/packages/mcp/src/keycardai/mcp/client/auth/transports.py index 67562a74..699f3ef9 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/transports.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/transports.py @@ -3,7 +3,7 @@ from collections.abc import AsyncGenerator from typing import TYPE_CHECKING -from httpx import Auth, Request, Response +from httpx2 import Auth, Request, Response from ..logging_config import get_logger @@ -15,9 +15,9 @@ class HttpxAuth(Auth): """ - Adapts AuthStrategy to httpx.Auth interface. + Adapts AuthStrategy to httpx2.Auth interface. - This is a thin adapter that translates between httpx's auth flow + This is a thin adapter that translates between httpx2's auth flow and our transport-agnostic AuthStrategy protocol. All business logic (OAuth discovery, token management, etc.) is @@ -38,7 +38,7 @@ def __init__(self, strategy: "AuthStrategy"): async def async_auth_flow(self, request: Request) -> AsyncGenerator[Request, Response]: """ - httpx auth flow - delegates to AuthStrategy. + httpx2 auth flow - delegates to AuthStrategy. This method: 1. Adds auth metadata to the request (if available) @@ -90,4 +90,3 @@ async def async_auth_flow(self, request: Request) -> AsyncGenerator[Request, Res except Exception as e: logger.error(f"Error handling auth challenge: {e}", exc_info=True) - diff --git a/packages/mcp/src/keycardai/mcp/client/connection/http.py b/packages/mcp/src/keycardai/mcp/client/connection/http.py index 02976f8b..7087ea68 100644 --- a/packages/mcp/src/keycardai/mcp/client/connection/http.py +++ b/packages/mcp/src/keycardai/mcp/client/connection/http.py @@ -1,11 +1,13 @@ """ HTTP connection implementation for MCP client. -Provides StreamableHttpConnection which uses httpx for HTTP-based MCP connections. +Provides StreamableHttpConnection which uses httpx2 for HTTP-based MCP connections. """ +from contextlib import AbstractAsyncContextManager from typing import TYPE_CHECKING, Any -from mcp.client.streamable_http import streamable_http_client +import httpx2 +from mcp.client.streamable_http import TransportStreams, streamable_http_client from ..auth.strategies import create_auth_strategy from ..auth.transports import HttpxAuth @@ -19,6 +21,9 @@ logger = get_logger(__name__) +MCP_DEFAULT_TIMEOUT = 30.0 +MCP_DEFAULT_SSE_READ_TIMEOUT = 300.0 + class StreamableHttpConnection(Connection): """ @@ -52,7 +57,10 @@ def __init__( self.server_config = server_config self.context = context self.coordinator = coordinator - self._mcp_client = None + self._mcp_client: ( + AbstractAsyncContextManager[TransportStreams] | None + ) = None + self._http_client: httpx2.AsyncClient | None = None self._disconnecting = False # Create connection-specific sub-namespace @@ -76,31 +84,57 @@ def __init__( async def connect(self) -> tuple[Any, Any]: """Establish HTTP connection with auth adapter.""" - # Create httpx auth adapter - # Strategy already has its storage from constructor + url = self.server_config.get("url") + if not isinstance(url, str): + raise ValueError("HTTP server configuration requires a URL") + auth = HttpxAuth(strategy=self.auth_strategy) + self._http_client = httpx2.AsyncClient( + auth=auth, + follow_redirects=True, + timeout=httpx2.Timeout( + MCP_DEFAULT_TIMEOUT, + read=MCP_DEFAULT_SSE_READ_TIMEOUT, + ), + ) self._mcp_client = streamable_http_client( - self.server_config.get("url"), - auth=auth, + url, + http_client=self._http_client, ) - self._read_stream, self._write_stream, _ = await self._mcp_client.__aenter__() - return (self._read_stream, self._write_stream) + try: + self._read_stream, self._write_stream = ( + await self._mcp_client.__aenter__() + ) + return (self._read_stream, self._write_stream) + except Exception: + await self.disconnect() + raise async def disconnect(self) -> None: """Disconnect from HTTP server.""" - if self._disconnecting or self._mcp_client is None: + if self._disconnecting or ( + self._mcp_client is None and self._http_client is None + ): return self._disconnecting = True try: - await self._mcp_client.__aexit__(None, None, None) - except Exception as e: - # Log but don't raise - we're cleaning up and want to ensure - # resources are released even if there are background task errors - logger.debug(f"Error during disconnect (suppressed): {e}") + if self._mcp_client is not None: + try: + await self._mcp_client.__aexit__(None, None, None) + except Exception as e: + logger.debug(f"Error during disconnect (suppressed): {e}") + + if self._http_client is not None: + try: + await self._http_client.aclose() + except Exception as e: + logger.debug( + f"Error closing HTTP client (suppressed): {e}" + ) finally: self._mcp_client = None + self._http_client = None self._disconnecting = False - diff --git a/packages/mcp/src/keycardai/mcp/client/session.py b/packages/mcp/src/keycardai/mcp/client/session.py index a1aae5d4..2e2338f5 100644 --- a/packages/mcp/src/keycardai/mcp/client/session.py +++ b/packages/mcp/src/keycardai/mcp/client/session.py @@ -377,7 +377,10 @@ async def _handle_connection_failure(self, error: Exception) -> None: Args: error: The exception that occurred during connection """ - logger.error("Failed to establish connection") + logger.error( + f"Failed to establish connection: {error}", + exc_info=(type(error), error, error.__traceback__), + ) await self._cleanup_failed_connection() error_status = self._classify_connection_error(error) diff --git a/packages/mcp/tests/integration/test_streamable_http_client.py b/packages/mcp/tests/integration/test_streamable_http_client.py new file mode 100644 index 00000000..b77e1587 --- /dev/null +++ b/packages/mcp/tests/integration/test_streamable_http_client.py @@ -0,0 +1,139 @@ +"""Integration coverage for the real MCP streamable HTTP transport.""" + +import asyncio +import contextlib +import socket +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +import pytest +import uvicorn +from mcp.server.mcpserver import MCPServer +from starlette.responses import Response +from starlette.types import ASGIApp, Receive, Scope, Send + +from keycardai.mcp.client import Client +from keycardai.mcp.client.connection.http import StreamableHttpConnection +from keycardai.mcp.client.session import SessionStatus + +API_KEY = "integration-test-key" + + +class ApiKeyMiddleware: + """Require an API key for HTTP requests while forwarding ASGI lifespan.""" + + def __init__(self, app: ASGIApp): + self.app = app + + async def __call__( + self, + scope: Scope, + receive: Receive, + send: Send, + ) -> None: + if scope["type"] == "http": + headers = dict(scope["headers"]) + if headers.get(b"x-api-key") != API_KEY.encode(): + await Response(status_code=401)(scope, receive, send) + return + + await self.app(scope, receive, send) + + +@asynccontextmanager +async def run_streamable_http_server( + require_api_key: bool, +) -> AsyncIterator[str]: + """Serve a minimal MCP app on an ephemeral localhost port.""" + mcp = MCPServer("streamable-http-integration") + + @mcp.tool() + def ping() -> str: + return "pong" + + app: ASGIApp = mcp.streamable_http_app() + if require_api_key: + app = ApiKeyMiddleware(app) + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + sock.listen(2048) + host, port = sock.getsockname() + + server = uvicorn.Server( + uvicorn.Config( + app, + log_level="warning", + lifespan="on", + ) + ) + server_task = asyncio.create_task(server.serve(sockets=[sock])) + + async def wait_until_started() -> None: + while not server.started: + if server_task.done(): + await server_task + await asyncio.sleep(0.01) + + try: + await asyncio.wait_for(wait_until_started(), timeout=5) + yield f"http://{host}:{port}/mcp" + finally: + server.should_exit = True + await asyncio.wait_for(server_task, timeout=5) + with contextlib.suppress(OSError): + sock.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("auth_config", "require_api_key"), + [ + (None, False), + ( + { + "type": "api_key", + "key": API_KEY, + "header_name": "X-API-Key", + }, + True, + ), + ], + ids=["unauthenticated", "api-key"], +) +async def test_client_connects_over_real_streamable_http_transport( + auth_config: dict[str, Any] | None, + require_api_key: bool, +) -> None: + """Client.connect reaches an operational session through the real transport.""" + async with run_streamable_http_server(require_api_key) as url: + server_config: dict[str, Any] = {"url": url} + if auth_config is not None: + server_config["auth"] = auth_config + + client = Client({"test": server_config}) + try: + await client.connect() + + session = client.sessions["test"] + assert session.status is SessionStatus.CONNECTED + assert session.is_operational + + connection = session._connection + assert isinstance(connection, StreamableHttpConnection) + http_client = connection._http_client + assert http_client is not None + assert http_client.follow_redirects is True + assert http_client.timeout.connect == 30.0 + assert http_client.timeout.read == 300.0 + assert http_client.timeout.write == 30.0 + assert http_client.timeout.pool == 30.0 + + tools = await client.list_tools("test") + assert [tool.tool.name for tool in tools] == ["ping"] + finally: + await client.disconnect() + + assert http_client.is_closed diff --git a/packages/mcp/tests/keycardai/mcp/client/auth/test_transports.py b/packages/mcp/tests/keycardai/mcp/client/auth/test_transports.py new file mode 100644 index 00000000..49cafc32 --- /dev/null +++ b/packages/mcp/tests/keycardai/mcp/client/auth/test_transports.py @@ -0,0 +1,80 @@ +"""Tests for HTTP authentication transport adapters.""" + +from typing import Any + +import httpx2 +import pytest + +from keycardai.mcp.client.auth.transports import HttpxAuth + + +class StubAuthStrategy: + """Authentication strategy stub for exercising the httpx2 auth flow.""" + + def __init__(self, metadata: list[dict[str, Any]], retry: bool = False): + self.metadata = metadata + self.retry = retry + self.challenges: list[httpx2.Response] = [] + + async def get_auth_metadata(self) -> dict[str, Any]: + return self.metadata.pop(0) + + async def handle_challenge( + self, + challenge: httpx2.Response, + resource_url: str, + ) -> bool: + self.challenges.append(challenge) + assert resource_url == "https://example.com/mcp" + return self.retry + + +@pytest.mark.asyncio +async def test_httpx_auth_adds_strategy_headers() -> None: + """The adapter applies strategy metadata through httpx2.""" + strategy = StubAuthStrategy([{"headers": {"X-API-Key": "secret"}}]) + + async def handle(request: httpx2.Request) -> httpx2.Response: + assert request.headers["X-API-Key"] == "secret" + return httpx2.Response(200) + + async with httpx2.AsyncClient( + auth=HttpxAuth(strategy), + transport=httpx2.MockTransport(handle), + ) as client: + response = await client.get("https://example.com/mcp") + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_httpx_auth_handles_challenge_and_retries() -> None: + """The adapter passes httpx2 challenges to the strategy before retrying.""" + strategy = StubAuthStrategy( + [ + {}, + {"headers": {"Authorization": "Bearer refreshed"}}, + ], + retry=True, + ) + request_count = 0 + + async def handle(request: httpx2.Request) -> httpx2.Response: + nonlocal request_count + request_count += 1 + if request_count == 1: + return httpx2.Response(401) + + assert request.headers["Authorization"] == "Bearer refreshed" + return httpx2.Response(200) + + async with httpx2.AsyncClient( + auth=HttpxAuth(strategy), + transport=httpx2.MockTransport(handle), + ) as client: + response = await client.get("https://example.com/mcp") + + assert response.status_code == 200 + assert request_count == 2 + assert len(strategy.challenges) == 1 + assert isinstance(strategy.challenges[0], httpx2.Response) diff --git a/packages/mcp/tests/keycardai/mcp/client/test_session.py b/packages/mcp/tests/keycardai/mcp/client/test_session.py index b043ce21..2585312b 100644 --- a/packages/mcp/tests/keycardai/mcp/client/test_session.py +++ b/packages/mcp/tests/keycardai/mcp/client/test_session.py @@ -5,6 +5,7 @@ lifecycle transitions. """ +import logging from typing import Any from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -219,6 +220,44 @@ async def test_connect_creates_connection_and_session(self): assert mock_client_session.entered is True assert mock_client_session.initialize_called is True + @pytest.mark.asyncio + async def test_connect_logs_transport_type_error_with_traceback( + self, + caplog: pytest.LogCaptureFixture, + ): + """Programming errors retain their traceback while status handling stays intact.""" + storage = InMemoryBackend() + coordinator = MockAuthCoordinator(storage) + context = coordinator.create_context("user:alice") + session = Session( + "test_server", + {"url": "http://localhost:3000"}, + context, + coordinator, + ) + mock_connection = MockConnection() + mock_connection.should_raise_on_start = TypeError( + "streamable_http_client() got an unexpected keyword argument 'auth'" + ) + + with ( + patch( + "keycardai.mcp.client.session.create_connection", + return_value=mock_connection, + ), + caplog.at_level(logging.ERROR), + ): + await session.connect() + + failure_record = next( + record + for record in caplog.records + if "Failed to establish connection" in record.message + ) + assert failure_record.exc_info is not None + assert failure_record.exc_info[0] is TypeError + assert session.status is SessionStatus.AUTH_FAILED + @pytest.mark.asyncio async def test_connect_when_already_connected_returns_early(self): """Test that connect returns early when already connected.""" @@ -1639,4 +1678,3 @@ async def test_requires_auth_method_still_works(self): session.get_auth_challenge = AsyncMock(return_value={"state": "test"}) assert await session.requires_auth() - diff --git a/uv.lock b/uv.lock index 81df9e95..deff6ea7 100644 --- a/uv.lock +++ b/uv.lock @@ -1490,6 +1490,7 @@ dependencies = [ { name = "aiohttp" }, { name = "aiosqlite" }, { name = "httpx" }, + { name = "httpx2" }, { name = "keycardai-oauth" }, { name = "keycardai-starlette" }, { name = "mcp" }, @@ -1512,6 +1513,7 @@ requires-dist = [ { name = "aiohttp", specifier = ">=3.11.11" }, { name = "aiosqlite", specifier = ">=0.20.0" }, { name = "httpx", specifier = ">=0.27.2" }, + { name = "httpx2", specifier = ">=2.5.0" }, { name = "keycardai-oauth", editable = "packages/oauth" }, { name = "keycardai-starlette", editable = "packages/starlette" }, { name = "langchain", marker = "extra == 'test'", specifier = ">=1.0.5" },