trymemcp/app/client.py

"""Async HTTP client for the TryMe Gaming App API."""

from __future__ import annotations

from http.cookies import SimpleCookie
from typing import Any

import httpx
from fastmcp.server.dependencies import get_http_headers

from app.config import Settings, settings


class TryMeApiError(RuntimeError):
    """Raised when the upstream TryMe API returns an error response."""

    def __init__(self, status_code: int, detail: Any, path: str) -> None:
        self.status_code = status_code
        self.detail = detail
        self.path = path
        super().__init__(f"TryMe API error ({status_code}) on {path}: {detail}")


def _normalize_authorization(value: str) -> str:
    """Ensure raw tokens are prefixed with 'Bearer ' while preserving full auth headers."""
    cleaned = value.strip()
    if not cleaned:
        return ""
    if " " in cleaned:
        return cleaned
    return f"Bearer {cleaned}"


def _extract_cookie_from_header(cookie_header: str | None, name: str) -> str | None:
    """Extract a named cookie from an incoming HTTP Cookie header string."""
    if not cookie_header:
        return None
    cookie = SimpleCookie()
    try:
        cookie.load(cookie_header)
    except Exception:
        return None
    morsel = cookie.get(name)
    return morsel.value if morsel else None


def _safe_incoming_headers() -> dict[str, str]:
    """Read incoming MCP HTTP headers if running inside an HTTP request context."""
    try:
        raw = get_http_headers(include_all=True) or {}
        return {k.lower(): v for k, v in raw.items() if isinstance(v, str)}
    except Exception:
        return {}


class TryMeClient:
    """HTTP client wrapper for interacting with https://tryme.tpk.pw."""

    def __init__(
        self,
        config: Settings | None = None,
        transport: httpx.AsyncBaseTransport | None = None,
    ) -> None:
        self.config = config or settings
        self._transport = transport

    def build_auth(
        self,
        *,
        api_key: str | None = None,
        authorization: str | None = None,
        tryme_session: str | None = None,
    ) -> tuple[dict[str, str], dict[str, str]]:
        """Resolve authentication headers and cookies with precedence:
        1. Explicit tool parameters
        2. Incoming MCP HTTP request headers/cookies
        3. Server environment variables
        """
        incoming = _safe_incoming_headers()
        headers: dict[str, str] = {}
        cookies: dict[str, str] = {}

        resolved_api_key = (
            api_key
            or incoming.get("x-api-key")
            or self.config.tryme_api_key
        )
        if resolved_api_key:
            headers["X-API-Key"] = resolved_api_key.strip()

        resolved_auth = (
            authorization
            or incoming.get("authorization")
            or self.config.tryme_bearer_token
        )
        if resolved_auth:
            norm = _normalize_authorization(resolved_auth)
            if norm:
                headers["Authorization"] = norm

        resolved_session = (
            tryme_session
            or _extract_cookie_from_header(incoming.get("cookie"), "tryme_session")
            or self.config.tryme_session_cookie
        )
        if resolved_session:
            cookies["tryme_session"] = resolved_session.strip()

        return headers, cookies

    async def request(
        self,
        method: str,
        path: str,
        *,
        params: dict[str, Any] | None = None,
        json_body: Any = None,
        api_key: str | None = None,
        authorization: str | None = None,
        tryme_session: str | None = None,
        follow_redirects: bool = False,
    ) -> Any:
        """Send an async HTTP request to the TryMe API and parse the response."""
        headers, cookies = self.build_auth(
            api_key=api_key,
            authorization=authorization,
            tryme_session=tryme_session,
        )
        headers.setdefault("Accept", "application/json")

        cleaned_params: dict[str, Any] | None = None
        if params:
            cleaned_params = {k: v for k, v in params.items() if v is not None}

        async with httpx.AsyncClient(
            base_url=self.config.tryme_base_url,
            timeout=self.config.tryme_timeout_seconds,
            follow_redirects=follow_redirects,
            cookies=cookies,
            transport=self._transport,
        ) as http:
            response = await http.request(
                method=method.upper(),
                url=path,
                params=cleaned_params,
                json=json_body,
                headers=headers,
            )

        if 300 <= response.status_code < 400:
            result: dict[str, Any] = {
                "status_code": response.status_code,
                "redirect_url": response.headers.get("location", ""),
            }
            set_cookie = response.headers.get("set-cookie")
            if set_cookie:
                result["cookies_set"] = True
            return result

        if response.status_code == 204 or not response.content:
            return None

        if response.status_code >= 400:
            try:
                err_payload = response.json()
                detail = err_payload.get("detail", err_payload) if isinstance(err_payload, dict) else err_payload
            except Exception:
                detail = response.text
            raise TryMeApiError(response.status_code, detail, path)

        content_type = response.headers.get("content-type", "")
        if "application/json" in content_type:
            return response.json()

        try:
            return response.json()
        except Exception:
            return {"status_code": response.status_code, "text": response.text}