infracloud/mcp/server.py

from __future__ import annotations

import os
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from mcp.server.fastmcp import FastMCP
import psycopg
from psycopg.rows import dict_row

# Mapeamento de Projetos
PROJECTS = {
    "infracloud": Path("/root/infracloud"),
    "gohorsejobs": Path("/root/gohorsejobs"),
    "core": Path("/root/core"),
    "saveinmed": Path("/root/saveinmed"),
}

DEFAULT_PROJECT = "infracloud"

POSTGRES_DSN_ENV = "INFRA_MCP_POSTGRES_DSN"
TRANSPORT_ENV = "INFRA_MCP_TRANSPORT"
HOST_ENV = "INFRA_MCP_HOST"
PORT_ENV = "INFRA_MCP_PORT"

READ_ONLY_SCRIPT_PREFIXES = (
    "check_", "fetch_", "get_", "inspect_", "verify_", "final_status", "watch_",
)
MUTATING_SCRIPT_PREFIXES = (
    "approve_", "complete_", "fix_", "merge_", "retrigger_", "revert_",
)

mcp = FastMCP(
    "infracloud-multi-project",
    instructions=(
        "Source of truth for infracloud, gohorsejobs, and core repositories. "
        "Always specify the project parameter when working outside infracloud. "
        "Use real repository files and inventory markdown as reference."
    ),
    host=os.getenv(HOST_ENV, "127.0.0.1"),
    port=int(os.getenv(PORT_ENV, "8000")),
)

@dataclass(frozen=True)
class ScriptInfo:
    path: Path
    relative_path: str
    is_read_only: bool
    kind: str

def _get_project_root(project: str | None = None) -> Path:
    name = project or DEFAULT_PROJECT
    if name not in PROJECTS:
        raise ValueError(f"Project '{name}' not found. Available: {', '.join(PROJECTS.keys())}")
    return PROJECTS[name]

def _ensure_in_project(path: Path, project: str | None = None) -> Path:
    root = _get_project_root(project)
    resolved = path.resolve()
    if root not in resolved.parents and resolved != root:
        raise ValueError(f"Path escapes project root: {path}")
    return resolved

def _postgres_dsn() -> str | None:
    return os.getenv(POSTGRES_DSN_ENV, "").strip() or None

def _get_pg_connection():
    dsn = _postgres_dsn()
    if not dsn:
        raise ValueError(f"{POSTGRES_DSN_ENV} is not configured")
    return psycopg.connect(dsn, row_factory=dict_row)

def _ensure_mcp_tables() -> None:
    if not _postgres_dsn(): return
    with _get_pg_connection() as conn:
        with conn.cursor() as cur:
            cur.execute("""
                CREATE TABLE IF NOT EXISTS infra_mcp_notes (
                    id BIGSERIAL PRIMARY KEY,
                    scope TEXT NOT NULL,
                    project TEXT NOT NULL DEFAULT 'infracloud',
                    title TEXT NOT NULL,
                    body TEXT NOT NULL,
                    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
                )
            """)
        conn.commit()

@mcp.tool()
def list_repo_scripts(project: str | None = None, name_filter: str | None = None) -> list[dict[str, Any]]:
    """List scripts in scripts/auto-organized for the specified project."""
    root = _get_project_root(project)
    scripts_dir = root / "scripts" / "auto-organized"
    if not scripts_dir.exists(): return []

    results = []
    for path in sorted(scripts_dir.rglob("*")):
        if not path.is_file(): continue
        results.append({
            "name": path.name,
            "relative_path": path.relative_to(root).as_posix(),
            "read_only": path.name.lower().startswith(READ_ONLY_SCRIPT_PREFIXES),
            "kind": "shell" if path.suffix == ".sh" else "powershell" if path.suffix == ".ps1" else "other"
        })
    if name_filter:
        results = [s for s in results if name_filter.lower() in s["relative_path"].lower()]
    return results

@mcp.tool()
def grep_project(query: str, project: str | None = None, glob: str | None = None) -> dict[str, Any]:
    """Search the specified project for infrastructure or code terms."""
    root = _get_project_root(project)
    command = ["rg", "-n", query, str(root)]
    if glob: command.extend(["-g", glob])

    completed = subprocess.run(command, capture_output=True, text=True, timeout=30)
    results = completed.stdout.splitlines()
    return {
        "project": project or DEFAULT_PROJECT,
        "matches": results[:200],
        "truncated": len(results) > 200,
        "stderr": completed.stderr[-4000:],
    }

@mcp.tool()
def repo_layout_summary(project: str | None = None) -> dict[str, Any]:
    """Return a compact summary of the specified project layout."""
    root = _get_project_root(project)
    return {
        "project": project or DEFAULT_PROJECT,
        "root": str(root),
        "dirs": sorted(path.name for path in root.iterdir() if path.is_dir() and not path.name.startswith(".")),
        "important_files": [f.name for f in root.glob("*.md")]
    }

@mcp.tool()
def read_project_file(relative_path: str, project: str | None = None) -> dict[str, str]:
    """Read a specific file from the selected project."""
    root = _get_project_root(project)
    path = _ensure_in_project(root / relative_path, project)
    return {
        "project": project or DEFAULT_PROJECT,
        "content": path.read_text(encoding="utf-8", errors="replace")[:20000]
    }

@mcp.tool()
def add_project_note(title: str, body: str, project: str | None = None, scope: str = "general") -> dict[str, Any]:
    """Store an operational note for a specific project in the database."""
    _ensure_mcp_tables()
    p_name = project or DEFAULT_PROJECT
    with _get_pg_connection() as conn:
        with conn.cursor() as cur:
            cur.execute("""
                INSERT INTO infra_mcp_notes (scope, project, title, body)
                VALUES (%s, %s, %s, %s)
                RETURNING id, project, title, body, created_at
            """, (scope, p_name, title, body))
            row = cur.fetchone()
        conn.commit()
    return row

if __name__ == "__main__":
    _ensure_mcp_tables()
    mcp.run(transport=os.getenv(TRANSPORT_ENV, "stdio"))