AegisGitea-MCP/src/aegis_gitea_mcp/server.py

"""Main MCP server implementation with FastAPI and SSE support."""

import logging
from typing import Any, Dict

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import ValidationError

from aegis_gitea_mcp.audit import get_audit_logger
from aegis_gitea_mcp.config import get_settings
from aegis_gitea_mcp.gitea_client import GiteaClient
from aegis_gitea_mcp.mcp_protocol import (
    AVAILABLE_TOOLS,
    MCPListToolsResponse,
    MCPToolCallRequest,
    MCPToolCallResponse,
    get_tool_by_name,
)
from aegis_gitea_mcp.tools.repository import (
    get_file_contents_tool,
    get_file_tree_tool,
    get_repository_info_tool,
    list_repositories_tool,
)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="AegisGitea MCP Server",
    description="Security-first MCP server for controlled AI access to self-hosted Gitea",
    version="0.1.0",
)

# Global settings and audit logger
settings = get_settings()
audit = get_audit_logger()


# Tool dispatcher mapping
TOOL_HANDLERS = {
    "list_repositories": list_repositories_tool,
    "get_repository_info": get_repository_info_tool,
    "get_file_tree": get_file_tree_tool,
    "get_file_contents": get_file_contents_tool,
}


@app.on_event("startup")
async def startup_event() -> None:
    """Initialize server on startup."""
    logger.info(f"Starting AegisGitea MCP Server on {settings.mcp_host}:{settings.mcp_port}")
    logger.info(f"Connected to Gitea instance: {settings.gitea_base_url}")
    logger.info(f"Audit logging enabled: {settings.audit_log_path}")
    
    # Test Gitea connection
    try:
        async with GiteaClient() as gitea:
            user = await gitea.get_current_user()
            logger.info(f"Authenticated as bot user: {user.get('login', 'unknown')}")
    except Exception as e:
        logger.error(f"Failed to connect to Gitea: {e}")
        raise


@app.on_event("shutdown")
async def shutdown_event() -> None:
    """Cleanup on server shutdown."""
    logger.info("Shutting down AegisGitea MCP Server")


@app.get("/")
async def root() -> Dict[str, Any]:
    """Root endpoint with server information."""
    return {
        "name": "AegisGitea MCP Server",
        "version": "0.1.0",
        "status": "running",
        "mcp_version": "1.0",
    }


@app.get("/health")
async def health() -> Dict[str, str]:
    """Health check endpoint."""
    return {"status": "healthy"}


@app.get("/mcp/tools")
async def list_tools() -> JSONResponse:
    """List all available MCP tools.

    Returns:
        JSON response with list of tool definitions
    """
    response = MCPListToolsResponse(tools=AVAILABLE_TOOLS)
    return JSONResponse(content=response.model_dump())


@app.post("/mcp/tool/call")
async def call_tool(request: MCPToolCallRequest) -> JSONResponse:
    """Execute an MCP tool call.

    Args:
        request: Tool call request with tool name and arguments

    Returns:
        JSON response with tool execution result
    """
    correlation_id = request.correlation_id or audit.log_tool_invocation(
        tool_name=request.tool,
        params=request.arguments,
    )

    try:
        # Validate tool exists
        tool_def = get_tool_by_name(request.tool)
        if not tool_def:
            error_msg = f"Tool '{request.tool}' not found"
            audit.log_tool_invocation(
                tool_name=request.tool,
                correlation_id=correlation_id,
                result_status="error",
                error=error_msg,
            )
            raise HTTPException(status_code=404, detail=error_msg)

        # Get tool handler
        handler = TOOL_HANDLERS.get(request.tool)
        if not handler:
            error_msg = f"Tool '{request.tool}' has no handler implementation"
            audit.log_tool_invocation(
                tool_name=request.tool,
                correlation_id=correlation_id,
                result_status="error",
                error=error_msg,
            )
            raise HTTPException(status_code=500, detail=error_msg)

        # Execute tool with Gitea client
        async with GiteaClient() as gitea:
            result = await handler(gitea, request.arguments)

        audit.log_tool_invocation(
            tool_name=request.tool,
            correlation_id=correlation_id,
            result_status="success",
        )

        response = MCPToolCallResponse(
            success=True,
            result=result,
            correlation_id=correlation_id,
        )
        return JSONResponse(content=response.model_dump())

    except ValidationError as e:
        error_msg = f"Invalid arguments: {str(e)}"
        audit.log_tool_invocation(
            tool_name=request.tool,
            correlation_id=correlation_id,
            result_status="error",
            error=error_msg,
        )
        raise HTTPException(status_code=400, detail=error_msg)

    except Exception as e:
        error_msg = str(e)
        audit.log_tool_invocation(
            tool_name=request.tool,
            correlation_id=correlation_id,
            result_status="error",
            error=error_msg,
        )
        response = MCPToolCallResponse(
            success=False,
            error=error_msg,
            correlation_id=correlation_id,
        )
        return JSONResponse(content=response.model_dump(), status_code=500)


@app.get("/mcp/sse")
async def sse_endpoint(request: Request) -> StreamingResponse:
    """Server-Sent Events endpoint for MCP protocol.

    This enables real-time communication with ChatGPT using SSE.

    Returns:
        Streaming SSE response
    """
    async def event_stream():
        """Generate SSE events."""
        # Send initial connection event
        yield f"data: {{'event': 'connected', 'server': 'AegisGitea MCP', 'version': '0.1.0'}}\n\n"

        # Keep connection alive
        try:
            while True:
                if await request.is_disconnected():
                    break
                
                # Heartbeat every 30 seconds
                yield f"data: {{'event': 'heartbeat'}}\n\n"
                
                # Wait for next heartbeat (in production, this would handle actual events)
                import asyncio
                await asyncio.sleep(30)
                
        except Exception as e:
            logger.error(f"SSE stream error: {e}")

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


def main() -> None:
    """Run the MCP server."""
    import uvicorn

    settings = get_settings()
    
    uvicorn.run(
        "aegis_gitea_mcp.server:app",
        host=settings.mcp_host,
        port=settings.mcp_port,
        log_level=settings.log_level.lower(),
        reload=False,
    )


if __name__ == "__main__":
    main()