lovdata-chat/session-manager/routes/proxy.py

import os
import time
from urllib.parse import urlparse

from fastapi import APIRouter, HTTPException, Request, Response
from starlette.responses import RedirectResponse, StreamingResponse
import httpx

from session_manager import session_manager
from http_pool import make_http_request, stream_http_request
from logging_config import (
    RequestContext,
    log_request,
    log_session_operation,
    log_security_event,
)

router = APIRouter(tags=["proxy"])

ALL_METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"]


def get_session_from_cookie(request: Request) -> str:
    session_id = request.cookies.get("lovdata_session")
    if not session_id:
        raise HTTPException(
            status_code=400,
            detail="No active session - please access via /c/{id} first",
        )
    return session_id


@router.get("/c/{session_id}")
@router.get("/c/{session_id}/{path:path}")
async def enter_session(request: Request, session_id: str, path: str = ""):
    """Entry point: set session cookie and redirect to root."""
    session = await session_manager.get_session(session_id)
    if not session or session.status != "running":
        raise HTTPException(
            status_code=404, detail="Session not found or not running"
        )
    resp = RedirectResponse(url="/", status_code=302)
    resp.set_cookie(
        key="lovdata_session",
        value=session_id,
        httponly=True,
        samesite="lax",
        max_age=86400,
    )
    return resp


def _is_sse_request(request: Request, path: str) -> bool:
    """Detect SSE requests by Accept header or path convention."""
    accept = request.headers.get("accept", "")
    if "text/event-stream" in accept:
        return True
    # OpenCode uses /global/event and /event paths for SSE
    if path == "event" or path.endswith("/event"):
        return True
    return False


@router.api_route("/{path:path}", methods=ALL_METHODS)
async def proxy_root_to_container(request: Request, path: str):
    """Catch-all: proxy everything to the container identified by cookie."""
    session_id = get_session_from_cookie(request)
    return await _proxy_to_container(request, session_id, path)


async def _proxy_to_container(request: Request, session_id: str, path: str):
    start_time = time.time()

    with RequestContext():
        session = await session_manager.get_session(session_id)
        if not session or session.status != "running":
            raise HTTPException(
                status_code=404, detail="Session not found or not running"
            )

    docker_host = os.getenv("DOCKER_HOST", "http://docker-daemon:2375")
    parsed = urlparse(docker_host)
    container_host = parsed.hostname or "docker-daemon"

    container_url = f"http://{container_host}:{session.port}"

    url = f"{container_url}/{path}"
    if request.url.query:
        url += f"?{request.url.query}"

    body = await request.body()

    headers = dict(request.headers)
    headers.pop("host", None)

    if session.auth_token:
        headers["Authorization"] = f"Bearer {session.auth_token}"
        headers["X-Session-Token"] = session.auth_token
        headers["X-Session-ID"] = session.session_id

    # --- SSE streaming path ---
    if _is_sse_request(request, path):
        return await _proxy_sse(request, session_id, path, url, headers, body, start_time)

    # --- Buffered path ---
    try:
        log_session_operation(
            session_id, "proxy_request", method=request.method, path=path
        )

        response = await make_http_request(
            method=request.method,
            url=url,
            headers=headers,
            content=body,
        )

        duration_ms = (time.time() - start_time) * 1000
        log_request(
            request.method,
            f"/{path}",
            response.status_code,
            duration_ms,
            session_id=session_id,
            operation="proxy_complete",
        )

        resp = Response(
            content=response.content,
            status_code=response.status_code,
            headers=dict(response.headers),
        )
        resp.set_cookie(
            key="lovdata_session",
            value=session_id,
            httponly=True,
            samesite="lax",
            max_age=86400,
        )
        return resp

    except httpx.TimeoutException as e:
        duration_ms = (time.time() - start_time) * 1000
        log_request(
            request.method, f"/{path}", 504, duration_ms,
            session_id=session_id, error="timeout",
        )
        raise HTTPException(
            status_code=504, detail="Request to session container timed out"
        )
    except httpx.RequestError as e:
        duration_ms = (time.time() - start_time) * 1000
        log_request(
            request.method, f"/{path}", 502, duration_ms,
            session_id=session_id, error=str(e),
        )
        raise HTTPException(
            status_code=502,
            detail=f"Failed to connect to session container: {str(e)}",
        )


async def _proxy_sse(
    request: Request,
    session_id: str,
    path: str,
    url: str,
    headers: dict,
    body: bytes,
    start_time: float,
):
    """Proxy an SSE event stream without buffering."""
    log_session_operation(
        session_id, "proxy_sse_stream", method=request.method, path=path
    )

    async def event_generator():
        try:
            async with stream_http_request(
                method=request.method,
                url=url,
                headers=headers,
                content=body,
            ) as upstream:
                async for chunk in upstream.aiter_bytes():
                    yield chunk
        except httpx.RequestError as e:
            log_security_event(
                "proxy_sse_error", "error",
                session_id=session_id, method=request.method,
                path=path, error=str(e),
            )
        finally:
            duration_ms = (time.time() - start_time) * 1000
            log_request(
                request.method, f"/{path}", 200, duration_ms,
                session_id=session_id, operation="proxy_sse_complete",
            )

    resp = StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )
    resp.set_cookie(
        key="lovdata_session",
        value=session_id,
        httponly=True,
        samesite="lax",
        max_age=86400,
    )
    return resp