proxy.py

"""
FleetMind MCP Authentication Proxy
Captures API keys from initial SSE connections and injects them into tool requests.

This proxy sits between MCP clients and the FastMCP server, solving the
multi-tenant authentication problem by:
1. Capturing api_key from initial /sse?api_key=xxx connection
2. Storing api_key mapped to session_id
3. Injecting api_key into subsequent /messages/?session_id=xxx requests

Architecture:
  MCP Client -> Proxy (port 7860) -> FastMCP (port 7861)
"""

import asyncio
import logging
import os
from aiohttp import web, ClientSession, ClientTimeout
from aiohttp.client_exceptions import ClientConnectionResetError
from urllib.parse import urlencode, parse_qs, urlparse, urlunparse
import sys

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Proxy configuration
# On HuggingFace, PORT env var is set to 7860
PROXY_PORT = int(os.getenv("PORT", 7860))  # Public-facing port
FASTMCP_PORT = 7861  # Internal FastMCP server port (fixed)
FASTMCP_HOST = "localhost"

# Session storage: session_id -> api_key
session_api_keys = {}


async def proxy_handler(request):
    """
    Main proxy handler - forwards all requests to FastMCP server.
    Captures API keys from SSE connections and injects them into tool calls.
    """
    path = request.path
    query_params = dict(request.query)

    # Extract API key if present (initial SSE connection)
    api_key = query_params.get('api_key')
    session_id = query_params.get('session_id')

    # STEP 1: Capture API key from initial SSE connection
    if api_key and path == '/sse':
        logger.info(f"[AUTH] Captured API key from SSE connection: {api_key[:20]}...")
        # Store temporarily - will be linked to session when we see it
        session_api_keys['_pending_api_key'] = api_key

    # STEP 2: Link session_id to API key (from /messages requests)
    if session_id and path.startswith('/messages'):
        # Check if we have a stored API key for this session
        if session_id not in session_api_keys:
            # Link this session to the pending API key
            if '_pending_api_key' in session_api_keys:
                api_key_to_store = session_api_keys['_pending_api_key']
                session_api_keys[session_id] = api_key_to_store
                logger.info(f"[AUTH] Linked session {session_id[:12]}... to API key")

        # STEP 3: Inject API key into request for FastMCP
        stored_api_key = session_api_keys.get(session_id)
        if stored_api_key:
            query_params['api_key'] = stored_api_key
            logger.debug(f"[AUTH] Injected API key into request for session {session_id[:12]}...")

    # Build target URL for FastMCP server
    query_string = urlencode(query_params) if query_params else ""
    target_url = f"http://{FASTMCP_HOST}:{FASTMCP_PORT}{path}"
    if query_string:
        target_url += f"?{query_string}"

    # Forward request to FastMCP
    # For SSE connections: total=None disables overall timeout (keeps connection alive)
    # Still use socket timeouts for safety (sock_connect, sock_read)
    async with ClientSession(
        timeout=ClientTimeout(
            total=None,         # No total timeout for long-lived SSE connections
            sock_connect=30,    # 30 seconds for initial connection
            sock_read=300       # 5 minutes for individual socket reads
        )
    ) as session:
        try:
            # Copy headers
            headers = dict(request.headers)
            # Remove host header to avoid conflicts
            headers.pop('Host', None)

            # Forward request based on method
            if request.method == 'GET':
                async with session.get(target_url, headers=headers) as resp:
                    # For SSE, stream the response
                    if 'text/event-stream' in resp.content_type:
                        # Create streaming response for SSE
                        response = web.StreamResponse(
                            status=resp.status,
                            reason=resp.reason,
                            headers=dict(resp.headers)
                        )
                        await response.prepare(request)

                        # Background task to send keep-alive pings (prevents timeout)
                        async def send_keepalive():
                            try:
                                while True:
                                    await asyncio.sleep(30)  # Send ping every 30 seconds
                                    await response.write(b":\n\n")  # SSE comment (ignored by client)
                            except asyncio.CancelledError:
                                pass

                        keepalive_task = asyncio.create_task(send_keepalive())

                        try:
                            # Stream chunks from FastMCP to client
                            async for chunk in resp.content.iter_any():
                                await response.write(chunk)

                            await response.write_eof()
                        finally:
                            # Cancel keep-alive task when streaming completes
                            keepalive_task.cancel()
                            try:
                                await keepalive_task
                            except asyncio.CancelledError:
                                pass

                        return response
                    else:
                        # For regular responses, read entire body
                        body = await resp.read()
                        resp_headers = dict(resp.headers)
                        return web.Response(
                            body=body,
                            status=resp.status,
                            headers=resp_headers
                        )

            elif request.method == 'POST':
                body = await request.read()
                async with session.post(target_url, data=body, headers=headers) as resp:
                    resp_body = await resp.read()
                    # Don't pass content_type separately - it's already in headers
                    resp_headers = dict(resp.headers)
                    return web.Response(
                        body=resp_body,
                        status=resp.status,
                        headers=resp_headers
                    )

            else:
                # Forward other methods
                async with session.request(
                    request.method,
                    target_url,
                    data=await request.read(),
                    headers=headers
                ) as resp:
                    body = await resp.read()
                    return web.Response(
                        body=body,
                        status=resp.status,
                        headers=dict(resp.headers)
                    )

        except (ClientConnectionResetError, ConnectionResetError) as e:
            # Client disconnected - this is normal for SSE connections
            # Log at DEBUG level to reduce noise
            logger.debug(f"[SSE] Client disconnected: {e}")
            return web.Response(text="Client disconnected", status=499)

        except Exception as e:
            import traceback
            error_details = traceback.format_exc()
            logger.error(f"[ERROR] Proxy error: {type(e).__name__}: {e}")
            logger.error(f"[ERROR] Traceback:\n{error_details}")
            return web.Response(
                text=f"Proxy error: {type(e).__name__}: {str(e)}",
                status=502
            )


async def health_check(request):
    """Health check endpoint"""
    return web.Response(text="FleetMind Proxy OK", status=200)


def create_app():
    """Create and configure the proxy application"""
    app = web.Application()

    # Health check endpoint
    app.router.add_get('/health', health_check)

    # Proxy all other requests
    app.router.add_route('*', '/{path:.*}', proxy_handler)

    return app


async def main():
    """Start the proxy server"""
    print("\n" + "=" * 70)
    print("FleetMind MCP Authentication Proxy")
    print("=" * 70)
    print(f"Proxy listening on: http://0.0.0.0:{PROXY_PORT}")
    print(f"Forwarding to FastMCP: http://{FASTMCP_HOST}:{FASTMCP_PORT}")
    print("=" * 70)
    print("[OK] Multi-tenant authentication enabled")
    print("[OK] API keys captured from SSE connections")
    print("[OK] Sessions automatically linked to API keys")
    print("=" * 70 + "\n")

    app = create_app()
    runner = web.AppRunner(app)
    await runner.setup()

    site = web.TCPSite(runner, '0.0.0.0', PROXY_PORT)
    await site.start()

    logger.info(f"[OK] Proxy server started on port {PROXY_PORT}")
    logger.info(f"[OK] Forwarding to FastMCP on {FASTMCP_HOST}:{FASTMCP_PORT}")

    # Keep running
    try:
        await asyncio.Event().wait()
    except KeyboardInterrupt:
        logger.info("Shutting down proxy server...")
        await runner.cleanup()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nProxy server stopped.")
        sys.exit(0)