proxy_server.py

# HF Spaces requires all traffic through port 7860
# This proxy routes:
#   /chat/* → Agent Zero (port 80)
#   /sandbox/* → Marimo (port 8080)
#   /api/* → OmniRoute (port 20128)

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
import httpx
import asyncio

app = FastAPI()

# Service ports (internal)
AGENT_ZERO_PORT = 80
MARIMO_PORT = 8080
OMNIROUTE_PORT = 20128
HF_PORT = 7860

# HTTP client with connection pooling
http_client = httpx.AsyncClient(timeout=60.0)


@app.get("/health")
async def health():
    """HF Spaces healthcheck endpoint"""
    return {"status": "healthy", "service": "marimo-zero"}


@app.get("/")
async def root():
    """Serve main UI"""
    try:
        with open("/app/static/index.html", "r") as f:
            return HTMLResponse(content=f.read())
    except FileNotFoundError:
        return HTMLResponse("<h1>Marimo Zero</h1><p>UI not found</p>", status_code=404)


@app.api_route("/chat/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy_agent_zero(request: Request, path: str):
    """Proxy to Agent Zero"""
    return await proxy_request(AGENT_ZERO_PORT, path, request)


@app.api_route("/sandbox/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy_marimo(request: Request, path: str):
    """Proxy to Marimo"""
    return await proxy_request(MARIMO_PORT, path, request)


@app.api_route("/api/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy_omniroute(request: Request, path: str):
    """Proxy to OmniRoute"""
    return await proxy_request(OMNIROUTE_PORT, path, request)


async def proxy_request(target_port: int, path: str, request: Request):
    """Generic proxy handler"""
    url = f"http://localhost:{target_port}/{path}"
    
    try:
        # Forward headers
        headers = dict(request.headers)
        headers.pop("host", None)
        headers.pop("content-length", None)
        
        # Get body
        body = await request.body() if request.method in ["POST", "PUT"] else None
        
        # Make request
        response = await http_client.request(
            method=request.method,
            url=url,
            headers=headers,
            content=body
        )
        
        # Return response
        return StreamingResponse(
            response.aiter_bytes(),
            status_code=response.status_code,
            headers=dict(response.headers)
        )
    except httpx.ConnectError as e:
        raise HTTPException(status_code=503, detail=f"Service unavailable (port {target_port})")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


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