echo_server.py

from mcp.server.fastmcp import FastMCP
from typing import Optional, Literal
import httpx
import os
import re
from contextvars import ContextVar

# Try to import TokenVerifier from fastmcp; fall back gracefully if path differs
try:
    from fastmcp.server.auth.verifier import TokenVerifier  # type: ignore
except Exception:  # pragma: no cover - fallback for alternate installs
    TokenVerifier = object  # type: ignore

# Context variable to store the verified bearer token per-request
_current_bearer_token: ContextVar[Optional[str]] = ContextVar("current_bearer_token", default=None)


class OpenWeatherTokenVerifier(TokenVerifier):  # type: ignore[misc]
    """Minimal verifier that accepts an OpenWeather API key as the bearer token.

    For safety, we perform a lightweight format check (32 hex chars typical for OpenWeather keys).
    On success, we stash the token in a contextvar and return claims including it.
    """

    _hex32 = re.compile(r"^[a-fA-F0-9]{32}$")

    def verify_token(self, token: str) -> dict:  # type: ignore[override]
        if not token:
            raise ValueError("Missing bearer token")
        # Basic format validation (doesn't call OpenWeather)
        if not self._hex32.match(token):
            # Allow override via env to support alternative key formats if needed
            allow_any = os.getenv("OPENWEATHER_ALLOW_ANY_TOKEN", "false").lower() in {"1", "true", "yes"}
            if not allow_any:
                raise ValueError("Invalid OpenWeather token format; expected 32 hex characters")
        _current_bearer_token.set(token)
        return {"auth_type": "openweather", "openweather_token": token}


mcp = FastMCP(name="OpenWeatherServer", stateless_http=True, auth=OpenWeatherTokenVerifier())

BASE_URL = "https://api.openweathermap.org"
DEFAULT_TIMEOUT_SECONDS = 15.0


def _require_token() -> str:
    # Priority: verified bearer token → env var → error
    token = _current_bearer_token.get()
    if token:
        return token
    env_token = os.getenv("OPENWEATHER_API_KEY")
    if env_token:
        return env_token
    raise ValueError(
        "OpenWeather token required. Provide as HTTP Authorization: Bearer <token> or set OPENWEATHER_API_KEY."
    )


def _get(path: str, params: dict) -> dict:
    try:
        with httpx.Client(timeout=httpx.Timeout(DEFAULT_TIMEOUT_SECONDS)) as client:
            resp = client.get(f"{BASE_URL}{path}", params=params)
            resp.raise_for_status()
            return resp.json()
    except httpx.HTTPStatusError as e:
        status = e.response.status_code if e.response is not None else None
        text = e.response.text if e.response is not None else str(e)
        raise ValueError(f"OpenWeather API error (status={status}): {text}")
    except Exception as e:
        raise RuntimeError(f"OpenWeather request failed: {e}")


@mcp.tool(description="Get current weather using One Call 3.0 (current segment). Token is taken from Authorization bearer or OPENWEATHER_API_KEY env.")
def current_weather(
    lat: float,
    lon: float,
    units: Literal["standard", "metric", "imperial"] = "metric",
    lang: str = "en",
) -> dict:
    token = _require_token()
    params = {
        "lat": lat,
        "lon": lon,
        "appid": token,
        "units": units,
        "lang": lang,
        "exclude": "minutely,hourly,daily,alerts",
    }
    data = _get("/data/3.0/onecall", params)
    return {
        "coord": {"lat": lat, "lon": lon},
        "units": units,
        "lang": lang,
        "current": data.get("current", data),
    }


@mcp.tool(description="Get One Call 3.0 data for coordinates. Use 'exclude' to omit segments (comma-separated). Token comes from bearer or env.")
def onecall(
    lat: float,
    lon: float,
    exclude: Optional[str] = None,
    units: Literal["standard", "metric", "imperial"] = "metric",
    lang: str = "en",
) -> dict:
    token = _require_token()
    params = {
        "lat": lat,
        "lon": lon,
        "appid": token,
        "units": units,
        "lang": lang,
    }
    if exclude:
        params["exclude"] = exclude
    return _get("/data/3.0/onecall", params)


@mcp.tool(description="Get 5-day/3-hour forecast by city name. Token comes from bearer or env.")
def forecast_city(
    city: str,
    units: Literal["standard", "metric", "imperial"] = "metric",
    lang: str = "en",
) -> dict:
    token = _require_token()
    params = {"q": city, "appid": token, "units": units, "lang": lang}
    return _get("/data/2.5/forecast", params)


@mcp.tool(description="Get air pollution data for coordinates. Token comes from bearer or env.")
def air_pollution(lat: float, lon: float) -> dict:
    token = _require_token()
    params = {"lat": lat, "lon": lon, "appid": token}
    return _get("/data/2.5/air_pollution", params)