app.py

import asyncio
import logging
import os
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI, Request
from sse_starlette.sse import EventSourceResponse
from mcp.server.sse import SseServerTransport

import gradio as gr
from delegation_mcp.server import DelegationMCPServer
from delegation_mcp.gradio_monitor import create_monitor_ui
from mcp.server.models import InitializationOptions


# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global server instance
mcp_server = None
transport = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage server lifecycle."""
    global mcp_server

    # Initialize MCP server
    logger.info("Initializing Delegation MCP Server...")
    mcp_server = DelegationMCPServer(enable_auto_discovery=True)

    # Initialize agent discovery
    if mcp_server.enable_auto_discovery:
        await mcp_server._discover_and_register_agents()

    logger.info("MCP Server ready!")

    yield

    # Cleanup
    logger.info("Shutting down Delegation MCP Server...")

app = FastAPI(lifespan=lifespan)

@app.get("/sse")
async def handle_sse(request: Request):
    """Handle SSE connections for MCP."""
    global mcp_server, transport
    
    transport = SseServerTransport("/messages")
    
    async def event_generator():
        async with mcp_server.server.run(
            transport.read_stream, 
            transport.write_stream, 
            InitializationOptions(
                server_name="delegation-mcp",
                server_version="0.3.0",
                capabilities=mcp_server.server.get_capabilities(
                    notification_options=None,
                    experimental_capabilities={
                        "tool_discovery": {},
                        "on_demand_loading": {},
                        "agent_discovery": {},
                    },
                ),
            )
        ) as stream:
            async for message in stream:
                yield message
                
    return EventSourceResponse(event_generator())

@app.post("/messages")
async def handle_messages(request: Request):
    """Handle incoming messages for MCP."""
    global transport
    if transport:
        return await transport.handle_post_message(request)
    return {"error": "No active transport"}

# Initialize a temporary server for the Gradio UI demo
# This is separate from the MCP server instance above
temp_server_for_ui = DelegationMCPServer(enable_auto_discovery=False)

# Mount Gradio app
logger.info("Mounting Gradio monitor...")
monitor_app = create_monitor_ui(demo_server=temp_server_for_ui)
app = gr.mount_gradio_app(app, monitor_app, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)