src/GLUE_Agent_mcp.py

"""
Model Context Protocol (MCP) for GLUE_Agent

GLUE_Agent provides comprehensive multi-omics data integration tools for single-cell RNA-seq and ATAC-seq analysis. This framework enables preprocessing, model training, and visualization of integrated multi-modal datasets.

This MCP Server contains tools extracted from the following tutorial files:
1. preprocessing
    - glue_preprocess_scrna: Preprocess scRNA-seq data with HVG selection, normalization, and PCA
    - glue_preprocess_scatac: Preprocess scATAC-seq data with LSI dimension reduction
    - glue_construct_regulatory_graph: Construct prior regulatory graph linking RNA and ATAC features
2. training
    - glue_configure_datasets: Configure RNA-seq and ATAC-seq datasets for GLUE model training
    - glue_train_model: Train GLUE model for multi-omics integration
    - glue_check_integration_consistency: Evaluate integration quality with consistency scores
    - glue_generate_embeddings: Generate cell and feature embeddings from trained GLUE model

Remote endpoint layout
----------------------
GET  /            — health check (HF Spaces liveness probe)
POST /upload      — upload a local file to /data/inputs on the server
GET  /outputs/... — download a file previously written to /data/outputs
/mcp              — Streamable HTTP MCP endpoint (used by Claude Code)
"""

import os
from pathlib import Path

from fastapi import FastAPI, HTTPException, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from fastmcp import FastMCP

from tools.preprocessing import preprocessing_mcp
from tools.training import training_mcp

# These env vars are set in the Dockerfile; both tool modules honour them too.
INPUT_DIR = Path(os.getenv("PREPROCESSING_INPUT_DIR", "/data/inputs"))
OUTPUT_DIR = Path(os.getenv("PREPROCESSING_OUTPUT_DIR", "/data/outputs"))

INPUT_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# --- MCP server -----------------------------------------------------------
mcp = FastMCP(name="GLUE_Agent")
mcp.mount(preprocessing_mcp)
mcp.mount(training_mcp)

# Build the MCP HTTP sub-app first so we can forward its lifespan.
mcp_http_app = mcp.http_app(path="/")

# --- Outer FastAPI app ----------------------------------------------------
# The MCP sub-app's lifespan initialises the StreamableHTTPSessionManager
# task group.  Mounted sub-apps don't get their lifespans invoked by
# Starlette, so we must propagate it to the parent app explicitly.
app = FastAPI(title="GLUE_Agent MCP Server", lifespan=mcp_http_app.lifespan)

_OAUTH_NOT_CONFIGURED = {
    "error": "unsupported_authorization_type",
    "error_description": "OAuth is not configured for this MCP server.",
}


@app.get("/")
async def health_check():
    """Liveness probe for HF Spaces; also confirms the MCP endpoint path."""
    return {"status": "ok", "mcp_endpoint": "/mcp"}


# Claude Code currently probes OAuth discovery/registration endpoints for remote
# HTTP MCP servers. If those endpoints return a generic 404 body (e.g.
# {"detail":"Not Found"}), Claude tries to parse it as an OAuth error payload and
# fails. Returning a well-formed OAuth error JSON avoids that failure mode.
@app.get("/.well-known/oauth-protected-resource", include_in_schema=False)
async def oauth_protected_resource():
    return JSONResponse(status_code=404, content=_OAUTH_NOT_CONFIGURED)


@app.get("/.well-known/oauth-authorization-server", include_in_schema=False)
async def oauth_authorization_server():
    return JSONResponse(status_code=404, content=_OAUTH_NOT_CONFIGURED)


@app.post("/register", include_in_schema=False)
async def oauth_dynamic_client_register():
    return JSONResponse(status_code=404, content=_OAUTH_NOT_CONFIGURED)


@app.post("/upload")
async def upload_file(file: UploadFile):
    """
    Upload a local file to the server's input directory.

    Usage (from launch_remote_mcp.sh CLAUDE.md template):
        curl -F "file=@/absolute/local/path" ${REMOTE_MCP_URL}/upload
    """
    dest = INPUT_DIR / file.filename
    dest.parent.mkdir(parents=True, exist_ok=True)
    content = await file.read()
    dest.write_bytes(content)
    return {"filename": file.filename, "path": str(dest), "size": len(content)}


@app.get("/outputs/{name:path}")
async def download_output(name: str):
    """
    Download a file previously written by an MCP tool.

    Usage (from launch_remote_mcp.sh CLAUDE.md template):
        wget ${REMOTE_MCP_URL}/outputs/<output_filename>
    """
    # Resolve and guard against path-traversal attacks
    file_path = (OUTPUT_DIR / name).resolve()
    if not str(file_path).startswith(str(OUTPUT_DIR.resolve())):
        raise HTTPException(status_code=400, detail="Invalid file path")
    if not file_path.is_file():
        raise HTTPException(status_code=404, detail=f"Output file not found: {name}")
    return FileResponse(str(file_path), filename=file_path.name)


app.mount("/mcp", mcp_http_app)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=int(os.getenv("PORT", 7860)),
    )