"""FastAPI reverse proxy for multiple Streamlit applications.""" import asyncio import contextlib from pathlib import Path import httpx import websockets from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse, StreamingResponse from proxy.app_discovery import discover_apps app = FastAPI(title="Streamline Gateway") # Determine apps directory - try local first, then Docker path apps_dir = Path(__file__).parent.parent / "apps" if not apps_dir.exists(): apps_dir = Path("/app/apps") # Discover apps once at startup and store in memory APP_MAP = discover_apps(apps_dir) # HTTP client for proxying requests http_client = httpx.AsyncClient(timeout=30.0) @app.on_event("startup") async def startup_event() -> None: """Log discovered apps on startup.""" print(f"Discovered {len(APP_MAP)} apps:") for name, metadata in APP_MAP.items(): print( f" - {metadata['emoji']} {metadata['display_name']} ({name}) -> http://localhost:{metadata['port']}" ) @app.on_event("shutdown") async def shutdown_event() -> None: """Close HTTP client on shutdown.""" await http_client.aclose() @app.get("/", response_class=HTMLResponse) async def read_root() -> HTMLResponse: """Serve the navigation page with links to all discovered apps.""" if not APP_MAP: body = """ Streamline Gateway

Welcome to the Streamline Gateway!

No apps found. Add apps to the /apps directory.

""" return HTMLResponse(content=body) links = "\n".join( [ f'

' f"{meta['emoji']} {meta['display_name']}: {meta['description']}

" for name, meta in APP_MAP.items() ] ) body = f""" Streamline Gateway

Welcome to the Streamline Gateway!

Available Apps:

{links} """ return HTMLResponse(content=body) @app.get("/health") async def health_check() -> dict[str, str]: """Health check endpoint.""" return {"status": "healthy", "apps": list(APP_MAP.keys())} @app.websocket("/{app_name}/{path:path}") async def websocket_proxy(websocket: WebSocket, app_name: str, path: str) -> None: """ Proxy WebSocket connections to apps. Handles all WebSocket paths for any app (Streamlit, FastAPI, etc.). """ if app_name not in APP_MAP: await websocket.close(code=1008) # Policy Violation return target_port = APP_MAP[app_name]["port"] target_url = f"ws://localhost:{target_port}/{path}" await websocket.accept() client_closed = False server_closed = False # Forward cookies and other headers to the backend additional_headers = [] if "cookie" in websocket.headers: additional_headers.append(("cookie", websocket.headers["cookie"])) if "user-agent" in websocket.headers: additional_headers.append(("user-agent", websocket.headers["user-agent"])) try: async with websockets.connect(target_url, additional_headers=additional_headers) as remote_ws: # type: ignore # Create tasks for bidirectional communication async def forward_to_remote() -> None: """Forward messages from client to Streamlit app.""" nonlocal client_closed try: while True: message = await websocket.receive() if message.get("type") == "websocket.disconnect": client_closed = True break if "text" in message: data = message["text"] await remote_ws.send(data) # type: ignore elif "bytes" in message: data = message["bytes"] await remote_ws.send(data) # type: ignore except WebSocketDisconnect: client_closed = True except Exception: client_closed = True async def forward_to_client() -> None: """Forward messages from Streamlit app to client.""" nonlocal server_closed, client_closed try: async for message in remote_ws: # type: ignore if client_closed: break try: if isinstance(message, str): await websocket.send_text(message) else: await websocket.send_bytes(message) except RuntimeError as e: if "after sending 'websocket.close'" in str(e): client_closed = True break raise except Exception: pass finally: server_closed = True # Run both directions concurrently await asyncio.gather( forward_to_remote(), forward_to_client(), return_exceptions=True, ) except Exception: pass finally: if not client_closed: with contextlib.suppress(RuntimeError): await websocket.close() @app.api_route( "/{app_name}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"], ) async def http_proxy(request: Request, app_name: str, path: str) -> StreamingResponse: """ Proxy HTTP requests to Streamlit apps. Handles all HTTP methods and streams responses back to the client. """ if app_name not in APP_MAP: return StreamingResponse( content=iter([b"App not found"]), status_code=404, media_type="text/plain", ) target_port = APP_MAP[app_name]["port"] target_url = f"http://localhost:{target_port}/{path}" # Build the proxied request headers = dict(request.headers) # Remove host header to avoid conflicts headers.pop("host", None) try: req = http_client.build_request( method=request.method, url=target_url, headers=headers, content=await request.body(), params=request.query_params, ) resp = await http_client.send(req, stream=True) # Stream the response back return StreamingResponse( content=resp.aiter_raw(), status_code=resp.status_code, headers=dict(resp.headers), ) except httpx.ConnectError: return StreamingResponse( content=iter([f"App '{app_name}' is not available".encode()]), status_code=503, media_type="text/plain", ) except Exception: return StreamingResponse( content=iter([b"Internal proxy error"]), status_code=500, media_type="text/plain", )