OSVauco/agents/tools/mcp_tools.py

"""
mcp_tools.py — MCP toolset integrations for OSVauco ADK agent.
Supports:
  - Google BigQuery MCP (via StreamableHTTP, uses ADC/OAuth)
  - Google Maps MCP (via StreamableHTTP, uses API key)
  - OPAX-MCP (Vauco intern CI/CD-kanal, FunctionTools over REST)
  - Local stdio MCP server (dev/test only)

Requires: google-adk >= 2.0.0, google-auth, httpx
"""

import os
import google.auth
import google.auth.transport.requests
from google.adk.tools.mcp_tool.mcp_toolset import (
    MCPToolset,
    StdioServerParameters,
    StreamableHTTPConnectionParams,
)

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
OPAX_MCP_URL = os.environ.get("MCP_SERVER_URL", "https://opax-mcp-zjbqp3prqq-uc.a.run.app")

# --- BigQuery MCP (Google managed, OAuth) ---
BIGQUERY_MCP_URL = f"https://bigquery.googleapis.com/mcp/projects/{PROJECT_ID}"

def get_bigquery_mcp_toolset() -> MCPToolset:
    """Connect to Google's managed BigQuery MCP server using ADC (no key file)."""
    credentials, _ = google.auth.default(
        scopes=["https://www.googleapis.com/auth/bigquery"]
    )
    credentials.refresh(google.auth.transport.requests.Request())
    return MCPToolset(
        connection_params=StreamableHTTPConnectionParams(
            url=BIGQUERY_MCP_URL,
            headers={
                "Authorization": f"Bearer {credentials.token}",
                "x-goog-user-project": PROJECT_ID,
            },
        )
    )

# --- Google Maps MCP (API key from env/Secret Manager) ---
MAPS_MCP_URL = "https://maps.googleapis.com/mcp"

def get_maps_mcp_toolset() -> MCPToolset:
    """Connect to Google's managed Maps MCP server using an API key."""
    maps_api_key = os.environ.get("MAPS_API_KEY", "")
    if not maps_api_key:
        raise ValueError("MAPS_API_KEY env var not set. Add to Secret Manager and load at startup.")
    return MCPToolset(
        connection_params=StreamableHTTPConnectionParams(
            url=MAPS_MCP_URL,
            headers={"X-Goog-Api-Key": maps_api_key},
        )
    )


# --- OPAX-MCP (Vauco intern CI/CD-kanal) via FunctionTools ---
# Bruker httpx over REST i stedet for MCPToolset siden opax-mcp er en
# vanlig FastAPI-server, ikke en MCP streamable-server med session-handshake.

def _opax_identity_token() -> str:
    """
    Hent Cloud Run identity token fra GCE metadata server.
    Kjører på Cloud Run / GCE — audience må matche opax-mcp sin URL.
    Fallback til ADC access token for lokal utvikling.
    """
    import httpx
    metadata_url = (
        "http://metadata.google.internal/computeMetadata/v1/instance"
        f"/service-accounts/default/identity?audience={OPAX_MCP_URL}&format=full"
    )
    try:
        resp = httpx.get(
            metadata_url,
            headers={"Metadata-Flavor": "Google"},
            timeout=5,
        )
        if resp.status_code == 200 and resp.text.strip():
            return resp.text.strip()
    except Exception:
        pass
    # Fallback: ADC access token (lokal dev)
    credentials, _ = google.auth.default()
    credentials.refresh(google.auth.transport.requests.Request())
    return credentials.token


def _opax_headers() -> dict:
    """Bygg auth-headers for opax-mcp: Cloud Run identity token + MCP-nøkkel."""
    return {
        "Authorization": f"Bearer {_opax_identity_token()}",
        "X-MCP-Key": os.environ.get("MCP_SECRET", ""),
        "Content-Type": "application/json",
    }


def push_static(file_path: str, content: str, content_type: str = "text/html; charset=utf-8") -> dict:
    """Last opp en statisk fil (HTML/CSS/JS) til opax.vauco.no via GCS. file_path er relativ sti, f.eks. static/jason.html."""
    import httpx
    resp = httpx.post(
        f"{OPAX_MCP_URL}/tools/push_static",
        headers=_opax_headers(),
        json={"file_path": file_path, "content": content, "content_type": content_type},
        timeout=30,
    )
    resp.raise_for_status()
    return resp.json()

def get_build_status() -> dict:
    """Hent status på siste Cloud Build-kjøring for OSVauco. Returnerer status, commit, varighet og logg-URL."""
    import httpx
    resp = httpx.get(
        f"{OPAX_MCP_URL}/tools/get_build_status",
        headers=_opax_headers(),
        timeout=15,
    )
    resp.raise_for_status()
    return resp.json()

def get_logs(lines: int = 50, severity: str = "DEFAULT") -> dict:
    """Hent siste Cloud Run-logger for osvauco-agent. severity: ERROR | WARNING | INFO | DEFAULT."""
    import httpx
    resp = httpx.get(
        f"{OPAX_MCP_URL}/tools/get_logs",
        headers=_opax_headers(),
        params={"lines": lines, "severity": severity},
        timeout=15,
    )
    resp.raise_for_status()
    return resp.json()

def deploy_service(branch: str = "main", trigger_id: str = "") -> dict:
    """Trigger Cloud Build for å bygge og deploye osvauco-agent. Krever HITL-godkjenning fra Chris."""
    import httpx
    resp = httpx.post(
        f"{OPAX_MCP_URL}/tools/deploy_service",
        headers=_opax_headers(),
        json={"branch": branch, "trigger_id": trigger_id},
        timeout=15,
    )
    resp.raise_for_status()
    return resp.json()

def get_opax_tools() -> list:
    """Returner liste med ADK FunctionTool-er for opax-mcp."""
    from google.adk.tools import FunctionTool
    return [
        FunctionTool(func=push_static),
        FunctionTool(func=get_build_status),
        FunctionTool(func=get_logs),
        FunctionTool(func=deploy_service),
    ]


# --- Local stdio MCP server (dev/test only) ---
def get_local_stdio_toolset(command: str, args: list[str]) -> MCPToolset:
    """Connect to a local stdio MCP server for development and testing."""
    return MCPToolset(
        connection_params=StdioServerParameters(
            command=command,
            args=args,
        )
    )