hermes-agent/tools/mcp_oauth_provider.py

"""Shared ``OAuthClientProvider`` customizations for Hermes MCP OAuth.

Two code paths build an SDK provider — ``tools.mcp_oauth.build_oauth_auth``
(legacy public API) and ``tools.mcp_oauth_manager.MCPOAuthManager`` — and both
need the same real-world fixes and the same config → constructor-kwargs
plumbing. This module holds that shared core once; the origin modules keep
their own subclass (logger name, disk-watch hooks) on top of it.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from tools.mcp_oauth import HermesTokenStorage

logger = logging.getLogger(__name__)


class HermesProviderMixin:
    """Token-endpoint fixes layered over the SDK's ``OAuthClientProvider``.

    - Supabase-style dynamic registration returns a ``client_secret`` but omits
      ``token_endpoint_auth_method``; the SDK then treats the client as public,
      omits the secret, and the token endpoint rejects the exchange (looping the
      browser authorization page). Coerce the in-memory client info to
      ``client_secret_post`` right before token and refresh requests.
    - ``token_user_agent`` (``oauth.user_agent``) is stamped onto token-endpoint
      requests only — some authorization servers and WAFs reject httpx's default.
    - Any 2xx token/refresh response is accepted, and token bodies never leak
      into exception text or log output.

    Must precede the SDK class in the MRO. Subclasses set ``_hermes_logger`` so
    warnings keep their origin module's logger name.
    """

    _hermes_logger: logging.Logger = logger

    def __init__(self, *args: Any, token_user_agent: str | None = None, **kwargs: Any):
        super().__init__(*args, **kwargs)
        self._hermes_token_user_agent = token_user_agent

    def _prepare_token_request(self, request):
        """Stamp the configured User-Agent onto a token/refresh request."""
        ua = getattr(self, "_hermes_token_user_agent", None)  # tests build via __new__
        if ua:
            request.headers["User-Agent"] = ua
        return request

    def _coerce_client_secret_post(self) -> None:
        """Same rule as ``HermesTokenStorage._coerce_secret_auth_method``, applied
        to the in-memory client info right before a token-endpoint request."""
        info = self.context.client_info
        if not info:
            return
        from mcp.shared.auth import OAuthClientInformationFull
        from tools.mcp_oauth import HermesTokenStorage

        data = info.model_dump(mode="json", exclude_none=True)
        if HermesTokenStorage._coerce_secret_auth_method(data):
            self.context.client_info = OAuthClientInformationFull.model_validate(data)

    async def _exchange_token_authorization_code(self, *args: Any, **kwargs: Any):
        self._coerce_client_secret_post()
        return self._prepare_token_request(await super()._exchange_token_authorization_code(*args, **kwargs))

    async def _refresh_token(self):
        self._coerce_client_secret_post()
        return self._prepare_token_request(await super()._refresh_token())

    async def _store_tokens(self, token_response) -> None:
        self.context.current_tokens = token_response
        self.context.update_token_expiry(token_response)
        await self.context.storage.set_tokens(token_response)

    async def _handle_token_response(self, response):
        """Accept any 2xx token response; never echo the body into errors."""
        from mcp.client.auth.oauth2 import OAuthTokenError

        if not (200 <= response.status_code < 300):
            raise OAuthTokenError(f"Token exchange failed ({response.status_code})")
        from httpx import HTTPError
        from mcp.client.auth.utils import handle_token_response_scopes

        try:
            token_response = await handle_token_response_scopes(response)
        except (HTTPError, OAuthTokenError):
            raise OAuthTokenError("Invalid token response") from None
        await self._store_tokens(token_response)

    async def _handle_refresh_response(self, response) -> bool:
        """Accept any 2xx refresh response; never log the body."""
        if not (200 <= response.status_code < 300):
            self._hermes_logger.warning("Token refresh failed: %s", response.status_code)
            self.context.clear_tokens()
            return False
        from httpx import HTTPError
        from mcp.shared.auth import OAuthToken
        from pydantic import ValidationError

        try:
            token_response = OAuthToken.model_validate_json(await response.aread())
        except (HTTPError, ValidationError):
            self._hermes_logger.warning("Invalid refresh response: %s", response.status_code)
            self.context.clear_tokens()
            return False
        await self._store_tokens(token_response)
        return True


def prepare_oauth_config(server_name: str, server_url: str, oauth_config: dict | None) -> tuple[dict, "HermesTokenStorage"]:
    """Copy the ``oauth:`` block, apply provider defaults, open its token storage.

    The copy matters: later steps record ``_resolved_port`` / ``_cimd_url`` in
    the dict, which must never leak back into the caller's config.
    """
    from tools import mcp_oauth as mo

    cfg = dict(oauth_config or {})
    mo.apply_oauth_provider_defaults(cfg, server_name=server_name, server_url=server_url)
    return cfg, mo.HermesTokenStorage(server_name)


def build_provider_kwargs(cfg: dict, storage: "HermesTokenStorage", *, ssh_proxy_hint: bool) -> dict[str, Any]:
    """Resolve the callback port and return the shared provider constructor kwargs.

    Runs the port → client-metadata → pre-registration sequence (order matters:
    metadata needs the resolved port, pre-registration needs the metadata).
    ``ssh_proxy_hint`` lets the redirect handler tailor its remote-session hint
    to a configured proxy ``redirect_uri``. Helpers are looked up on
    ``tools.mcp_oauth`` at call time so tests can patch them there.
    """
    from tools import mcp_oauth as mo

    port = mo._configure_callback_port(cfg, storage)
    client_metadata = mo._build_client_metadata(cfg)
    mo._maybe_preregister_client(storage, cfg, client_metadata)
    redirect_uri = (cfg.get("redirect_uri") or None) if ssh_proxy_hint else None
    return {
        "client_metadata": client_metadata,
        "storage": storage,
        "redirect_handler": mo._make_redirect_handler(port, redirect_uri=redirect_uri),
        # mcp 2.0 dropped OAuthClientProvider's own `timeout`; the configured
        # `oauth.timeout` bounds the callback waiter's poll loop instead.
        "callback_handler": mo._make_callback_waiter(
            port, cfg.get("_cimd_url"), timeout=float(cfg.get("timeout", 300))
        ),
        "token_user_agent": mo.token_request_user_agent(cfg),
        **mo.cimd_provider_kwargs(cfg),
    }