agent-platform/services/api-gateway/app/infrastructure/proxy.py

from dataclasses import dataclass
from typing import Literal

import httpx
from core_shared.observability import PARENT_SPAN_ID_HEADER, SPAN_ID_HEADER, TRACE_ID_HEADER
from core_shared.security import build_internal_service_headers
from fastapi import Request, Response
from fastapi.responses import StreamingResponse

from app.bootstrap.settings import ApiGatewaySettings
from app.infrastructure.audit import mark_gateway_target
from app.infrastructure.request_context import REQUEST_ID_HEADER, get_gateway_request_context
from app.schemas.gateway import DownstreamServiceHealth

ProxyServiceName = Literal[
    "session-service",
    "tool-service",
    "model-gateway-service",
    "model-provider-service",
    "code-runner-service",
    "agent-service",
    "memory-service",
    "team-service",
    "skill-service",
    "human-service",
    "knowledge-service",
    "event-service",
    "identity-service",
    "scheduler-service",
]


@dataclass(frozen=True)
class ProxyTarget:
    service_name: ProxyServiceName
    base_url: str
    path_prefix: str
    health_path: str


class ServiceProxy:
    def __init__(self, *, settings: ApiGatewaySettings, timeout_seconds: float) -> None:
        self.settings = settings
        self.timeout_seconds = timeout_seconds

    async def forward(
        self,
        *,
        request: Request,
        target: ProxyTarget,
        path: str) -> Response:
        target_url = build_target_url(target=target, path=path)
        mark_gateway_target(
            request,
            target_service=target.service_name,
            target_url=target_url)
        headers = build_forward_headers(request)
        request_context = get_gateway_request_context(request)
        headers[REQUEST_ID_HEADER] = request_context.request_id
        trace_id = getattr(request.state, "trace_id", None)
        span_id = getattr(request.state, "span_id", None)
        if isinstance(trace_id, str):
            headers[TRACE_ID_HEADER] = trace_id
        if isinstance(span_id, str):
            headers[PARENT_SPAN_ID_HEADER] = span_id
        headers.update(build_internal_service_headers(self.settings))
        body = await request.body()

        if _expects_stream(request):
            client = httpx.AsyncClient(timeout=self.timeout_seconds)
            request_builder = client.build_request(
                method=request.method,
                url=target_url,
                params=request.query_params,
                headers=headers,
                content=body)
            upstream_response = await client.send(request_builder, stream=True)
            return StreamingResponse(
                _stream_upstream_response(client, upstream_response),
                status_code=upstream_response.status_code,
                headers=build_response_headers(upstream_response),
                media_type=upstream_response.headers.get("content-type"))

        async with httpx.AsyncClient(timeout=self.timeout_seconds) as client:
            upstream_response = await client.request(
                method=request.method,
                url=target_url,
                params=request.query_params,
                headers=headers,
                content=body)

        return Response(
            content=upstream_response.content,
            status_code=upstream_response.status_code,
            headers=build_response_headers(upstream_response),
            media_type=upstream_response.headers.get("content-type"))

    async def check_health(self, target: ProxyTarget) -> DownstreamServiceHealth:
        health_url = f"{target.base_url.rstrip('/')}{target.health_path}"
        try:
            async with httpx.AsyncClient(timeout=self.timeout_seconds) as client:
                response = await client.get(
                    health_url,
                    headers=build_internal_service_headers(self.settings))
        except httpx.HTTPError as exc:
            return DownstreamServiceHealth(
                service=target.service_name,
                status="error",
                url=health_url,
                error_message=str(exc))

        return DownstreamServiceHealth(
            service=target.service_name,
            status="ok" if response.is_success else "error",
            url=health_url,
            status_code=response.status_code,
            error_message=None if response.is_success else response.text)


def build_target_url(*, target: ProxyTarget, path: str) -> str:
    normalized_path = path.strip("/")
    if normalized_path:
        return f"{target.base_url.rstrip('/')}{target.path_prefix}/{normalized_path}"
    return f"{target.base_url.rstrip('/')}{target.path_prefix}"


def build_forward_headers(request: Request) -> dict[str, str]:
    skipped_headers = {
        "host",
        "content-length",
        "connection",
        REQUEST_ID_HEADER,
        "x-internal-service-token",
        "x-internal-service-name",
        TRACE_ID_HEADER,
        SPAN_ID_HEADER,
        PARENT_SPAN_ID_HEADER,
    }
    return {
        key: value
        for key, value in request.headers.items()
        if key.lower() not in skipped_headers
    }


def build_response_headers(response: httpx.Response) -> dict[str, str]:
    skipped_headers = {"content-length", "transfer-encoding", "connection"}
    return {
        key: value
        for key, value in response.headers.items()
        if key.lower() not in skipped_headers
    }


def _expects_stream(request: Request) -> bool:
    accept_header = request.headers.get("accept", "")
    return "text/event-stream" in accept_header or request.url.path.endswith("stream")


async def _stream_upstream_response(
    client: httpx.AsyncClient,
    response: httpx.Response):
    try:
        async for chunk in response.aiter_raw():
            yield chunk
    finally:
        await response.aclose()
        await client.aclose()