_httpx_utils.py

"""Utilities for creating and using httpx2 AsyncClient instances in the MCP transports."""

from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import Any, Protocol

import httpx2

__all__ = ["create_mcp_http_client", "MCP_DEFAULT_TIMEOUT", "MCP_DEFAULT_SSE_READ_TIMEOUT"]

# Default MCP timeout configuration
MCP_DEFAULT_TIMEOUT = 30.0  # General operations (seconds)
MCP_DEFAULT_SSE_READ_TIMEOUT = 300.0  # SSE streams - 5 minutes (seconds)

# The headers httpx2.AsyncClient.sse() adds to an event-stream request.
_SSE_HEADERS = {"Accept": "text/event-stream", "Cache-Control": "no-store"}

# How many redirects one auth-flow request may follow within its origin (see RedirectAwareAuth).
_AUTH_REDIRECT_LIMIT = 5


class McpHttpClientFactory(Protocol):  # pragma: no branch
    def __call__(  # pragma: no branch
        self,
        headers: dict[str, str] | None = None,
        timeout: httpx2.Timeout | None = None,
        auth: httpx2.Auth | None = None,
    ) -> httpx2.AsyncClient: ...


def create_mcp_http_client(
    headers: dict[str, str] | None = None,
    timeout: httpx2.Timeout | None = None,
    auth: httpx2.Auth | None = None,
) -> httpx2.AsyncClient:
    """Create an httpx2 AsyncClient with the MCP transports' default timeouts.

    The client uses a 30-second timeout for connect/write/pool and a 300-second
    read timeout, because a server may hold a response stream open. Redirect
    following is left at the httpx2 default (off): the MCP transports follow
    redirects within the endpoint's origin themselves, see `stream_within_origin`.

    Args:
        headers: Optional headers to include with all requests.
        timeout: Request timeout as httpx2.Timeout object. Defaults to 30s for
            connect/write/pool and 300s for read (for long-lived SSE streams).
        auth: Optional authentication handler.

    Returns:
        Configured httpx2.AsyncClient instance.

    Note:
        The returned AsyncClient must be used as a context manager to ensure
        proper cleanup of connections.
    """
    if timeout is None:
        timeout = httpx2.Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT)
    kwargs: dict[str, Any] = {"timeout": timeout}
    if headers is not None:
        kwargs["headers"] = headers
    if auth is not None:  # pragma: no cover
        kwargs["auth"] = auth
    return httpx2.AsyncClient(**kwargs)


def _within_origin(url: httpx2.URL, location: httpx2.URL) -> bool:
    """Whether `location` is on `url`'s origin, or is its https upgrade on the default ports.

    httpx2 normalises a scheme's default port to None and lower-cases hosts, so
    plain tuple comparison is exact. The upgrade rule is the one httpx2 itself
    uses to decide a redirect has not left the origin (`_is_https_redirect`).
    """
    if (url.scheme, url.host, url.port) == (location.scheme, location.host, location.port):
        return True
    return (
        url.host == location.host
        and url.scheme == "http"
        and url.port is None
        and location.scheme == "https"
        and location.port is None
    )


def next_request_within_origin(response: httpx2.Response) -> httpx2.Request | None:
    """The request that follows `response`'s redirect, if it is one the MCP transports follow.

    That is when httpx2 built a next request for it (a redirect status with a
    Location), the next request keeps the method (307/308, or any redirect of a
    GET: httpx2 turns a POST into a body-less GET for 301/302/303, which would
    drop the message), its URL stays within the origin of the request just sent
    (same scheme, host and port, or http to https on the same host with default
    ports), and the Location does not bring userinfo of its own (which httpx2
    would otherwise send as Basic auth; userinfo the configured URL already had
    is kept by a relative Location and is fine). None for anything else,
    including a non-redirect.
    """
    next_request = response.next_request
    if next_request is None:
        return None
    sent = response.request
    if (
        next_request.method != sent.method
        or (next_request.url.userinfo and next_request.url.userinfo != sent.url.userinfo)
        or not _within_origin(sent.url, next_request.url)
    ):
        return None
    return next_request


@asynccontextmanager
async def stream_within_origin(
    client: httpx2.AsyncClient, method: str, url: httpx2.URL | str, **kwargs: Any
) -> AsyncGenerator[httpx2.Response]:
    """`client.stream(...)`, following redirects only while they stay within the request's origin.

    An MCP transport talks to one configured endpoint, and everything on a request
    (headers, auth, body) was configured for that endpoint. A redirect that
    `next_request_within_origin` accepts, such as a 307/308 trailing-slash
    normalisation, is followed, at most `client.max_redirects` times. Any other
    redirect (or one past that budget) is not followed: the redirect response
    itself is yielded, the way httpx2 hands one back when `follow_redirects` is
    off, and the caller treats it as the non-success it is. The client's own
    `follow_redirects` setting is not consulted. Requests an `httpx2.Auth` flow
    makes during the call are sent without following either; the SDK's OAuth
    providers apply the same rule to their own requests.
    """
    request = client.build_request(method, url, **kwargs)
    followed = 0
    while True:
        response = await client.send(request, stream=True, follow_redirects=False)
        next_request = next_request_within_origin(response)
        if next_request is None or followed == client.max_redirects:
            break
        try:
            # Drain the redirect body so the connection returns to the pool, as httpx2 does when it follows.
            await response.aread()
        finally:
            await response.aclose()
        request = next_request
        followed += 1
    try:
        yield response
    finally:
        await response.aclose()


async def request_within_origin(
    client: httpx2.AsyncClient, method: str, url: httpx2.URL | str, **kwargs: Any
) -> httpx2.Response:
    """`client.request(...)` with the redirect handling of `stream_within_origin`."""
    async with stream_within_origin(client, method, url, **kwargs) as response:
        await response.aread()
    return response


@asynccontextmanager
async def sse_within_origin(
    client: httpx2.AsyncClient,
    url: httpx2.URL | str,
    *,
    headers: dict[str, str] | None = None,
    max_event_size: int | None = 1024 * 1024,
) -> AsyncGenerator[httpx2.EventSource]:
    """`client.sse(url)` with the redirect handling of `stream_within_origin`."""
    merged = httpx2.Headers(_SSE_HEADERS)
    merged.update(headers or {})
    async with stream_within_origin(client, "GET", url, headers=merged) as response:
        yield httpx2.EventSource(response, max_event_size=max_event_size)


def redirect_location(response: httpx2.Response) -> httpx2.URL | None:
    """Where `response` redirects to, for use in a message: without userinfo, query or fragment,
    which can carry state that does not belong in an error or a log line. None if not a redirect."""
    if response.next_request is None:
        return None
    return response.next_request.url.copy_with(userinfo=b"", query=None, fragment=None)


def redirect_note(response: httpx2.Response) -> str:
    """A suffix naming the location of a redirect response that was not followed, else empty."""
    location = redirect_location(response)
    if location is None:
        return ""
    return f" (redirected to {location}; not followed)"


class RedirectAwareAuth(ABC, httpx2.Auth):
    """An `httpx2.Auth` whose own requests follow redirects the way MCP transport requests do.

    The transports send every request with redirect following off and follow a
    redirect themselves only within the endpoint's origin (`stream_within_origin`).
    httpx2 applies that per-request setting to the requests an auth flow makes
    too (metadata discovery, registration, token), so on their own those would
    follow nothing. Subclasses write their flow as `_auth_flow`; this class
    drives it and, for each request the flow makes other than the one being
    authenticated, follows a redirect that `next_request_within_origin` accepts,
    up to `_AUTH_REDIRECT_LIMIT` times. Any other redirect response is handed
    to the flow as it is.
    """

    @abstractmethod
    def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
        """The subclass's flow, written as `httpx2.Auth.async_auth_flow` otherwise would be."""

    async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
        flow = self._auth_flow(request)
        try:
            outgoing = await flow.__anext__()
            while True:
                response = yield outgoing
                if outgoing is not request:
                    for _ in range(_AUTH_REDIRECT_LIMIT):
                        follow = next_request_within_origin(response)
                        if follow is None:
                            break
                        response = yield follow
                outgoing = await flow.asend(response)
        except StopAsyncIteration:
            return
        finally:
            await flow.aclose()