Aether/apps/aether-gateway/src/handlers/proxy/finalize.rs

use crate::audit::emit_admin_audit;
use crate::constants::{
    CONTROL_ENDPOINT_SIGNATURE_HEADER, CONTROL_EXECUTION_RUNTIME_HEADER, CONTROL_REQUEST_ID_HEADER,
    CONTROL_ROUTE_CLASS_HEADER, CONTROL_ROUTE_FAMILY_HEADER, CONTROL_ROUTE_KIND_HEADER,
    DEPENDENCY_REASON_HEADER, EXECUTION_PATH_HEADER, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
    TRACE_ID_HEADER,
};
use crate::control::GatewayControlDecision;
use crate::control::GatewayPublicRequestContext;
use crate::middleware::{should_downgrade_access_log, RequestLogEmitted};
use crate::AppState;
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
use axum::body::{Body, Bytes};
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
use std::time::Instant;
use tracing::{info, trace, warn};

pub(super) fn request_wants_stream(
    request_context: &GatewayPublicRequestContext,
    body: &axum::body::Bytes,
) -> bool {
    if request_context
        .request_path
        .contains(":streamGenerateContent")
    {
        return true;
    }
    if !request_context
        .request_content_type
        .as_deref()
        .map(|value| value.to_ascii_lowercase().contains("application/json"))
        .unwrap_or(false)
        || body.is_empty()
    {
        return false;
    }
    serde_json::from_slice::<serde_json::Value>(body)
        .ok()
        .and_then(|value| value.get("stream").and_then(|stream| stream.as_bool()))
        .unwrap_or(false)
}

pub(super) fn finalize_gateway_response(
    state: &AppState,
    mut response: Response<Body>,
    trace_id: &str,
    remote_addr: &std::net::SocketAddr,
    method: &http::Method,
    path_and_query: &str,
    control_decision: Option<&GatewayControlDecision>,
    execution_path: &'static str,
    started_at: &Instant,
    request_permit: Option<AdmissionPermit>,
) -> Response<Body> {
    attach_control_decision_headers(&mut response, control_decision);
    if !response.headers().contains_key(TRACE_ID_HEADER) {
        response.headers_mut().insert(
            HeaderName::from_static(TRACE_ID_HEADER),
            HeaderValue::from_str(trace_id).expect("trace id should be a valid header value"),
        );
    }
    response.headers_mut().insert(
        HeaderName::from_static(EXECUTION_PATH_HEADER),
        HeaderValue::from_static(execution_path),
    );

    let elapsed_ms = started_at.elapsed().as_millis() as u64;
    let dependency_reason = response
        .headers()
        .get(DEPENDENCY_REASON_HEADER)
        .and_then(|value| value.to_str().ok())
        .unwrap_or("none")
        .to_string();
    let local_execution_runtime_miss_reason = response
        .headers()
        .get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
        .and_then(|value| value.to_str().ok())
        .unwrap_or("none")
        .to_string();
    let route_class = control_decision
        .and_then(|decision| decision.route_class.as_deref())
        .unwrap_or("passthrough");
    let request_id = response
        .headers()
        .get(CONTROL_REQUEST_ID_HEADER)
        .and_then(|value| value.to_str().ok())
        .filter(|value| !value.trim().is_empty())
        .unwrap_or("-")
        .to_string();
    let auth_context = control_decision.and_then(|decision| decision.auth_context.as_ref());
    let user_id = auth_context
        .map(|auth_context| auth_context.user_id.as_str())
        .unwrap_or("-");
    let api_key_id = auth_context
        .map(|auth_context| auth_context.api_key_id.as_str())
        .unwrap_or("-");
    let status_code = response.status().as_u16();
    emit_admin_audit(
        &mut response,
        trace_id,
        method,
        path_and_query,
        control_decision,
    );
    if response.status().is_server_error() {
        warn!(
            event_name = "http_request_failed",
            log_type = "access",
            status = "failed",
            status_code,
            trace_id = %trace_id,
            request_id,
            remote_addr = %remote_addr,
            method = %method,
            path = %path_and_query,
            user_id,
            api_key_id,
            route_class,
            execution_path,
            dependency_reason = dependency_reason.as_str(),
            local_execution_runtime_miss_reason = local_execution_runtime_miss_reason.as_str(),
            elapsed_ms,
            "gateway request failed"
        );
    } else if should_downgrade_access_log(method, path_and_query) {
        trace!(
            event_name = "http_request_completed",
            log_type = "access",
            status = "completed",
            status_code,
            trace_id = %trace_id,
            request_id,
            remote_addr = %remote_addr,
            method = %method,
            path = %path_and_query,
            user_id,
            api_key_id,
            route_class,
            execution_path,
            dependency_reason = dependency_reason.as_str(),
            local_execution_runtime_miss_reason = local_execution_runtime_miss_reason.as_str(),
            elapsed_ms,
            "gateway completed request"
        );
    } else {
        info!(
            event_name = "http_request_completed",
            log_type = "access",
            status = "completed",
            status_code,
            trace_id = %trace_id,
            request_id,
            remote_addr = %remote_addr,
            method = %method,
            path = %path_and_query,
            user_id,
            api_key_id,
            route_class,
            execution_path,
            dependency_reason = dependency_reason.as_str(),
            local_execution_runtime_miss_reason = local_execution_runtime_miss_reason.as_str(),
            elapsed_ms,
            "gateway completed request"
        );
    }
    response.extensions_mut().insert(RequestLogEmitted);

    maybe_hold_axum_response_permit(response, request_permit)
}

fn attach_control_decision_headers(
    response: &mut Response<Body>,
    control_decision: Option<&GatewayControlDecision>,
) {
    let Some(control_decision) = control_decision else {
        return;
    };
    if !response.headers().contains_key(CONTROL_ROUTE_CLASS_HEADER) {
        response.headers_mut().insert(
            HeaderName::from_static(CONTROL_ROUTE_CLASS_HEADER),
            HeaderValue::from_str(
                control_decision
                    .route_class
                    .as_deref()
                    .unwrap_or("passthrough"),
            )
            .expect("route class should be a valid header value"),
        );
    }
    if !response
        .headers()
        .contains_key(CONTROL_EXECUTION_RUNTIME_HEADER)
    {
        response.headers_mut().insert(
            HeaderName::from_static(CONTROL_EXECUTION_RUNTIME_HEADER),
            HeaderValue::from_static(if control_decision.is_execution_runtime_candidate() {
                "true"
            } else {
                "false"
            }),
        );
    }
    if let Some(route_family) = control_decision.route_family.as_deref() {
        if !response.headers().contains_key(CONTROL_ROUTE_FAMILY_HEADER) {
            response.headers_mut().insert(
                HeaderName::from_static(CONTROL_ROUTE_FAMILY_HEADER),
                HeaderValue::from_str(route_family)
                    .expect("route family should be a valid header value"),
            );
        }
    }
    if let Some(route_kind) = control_decision.route_kind.as_deref() {
        if !response.headers().contains_key(CONTROL_ROUTE_KIND_HEADER) {
            response.headers_mut().insert(
                HeaderName::from_static(CONTROL_ROUTE_KIND_HEADER),
                HeaderValue::from_str(route_kind)
                    .expect("route kind should be a valid header value"),
            );
        }
    }
    if let Some(endpoint_signature) = control_decision.auth_endpoint_signature.as_deref() {
        if !response
            .headers()
            .contains_key(CONTROL_ENDPOINT_SIGNATURE_HEADER)
        {
            response.headers_mut().insert(
                HeaderName::from_static(CONTROL_ENDPOINT_SIGNATURE_HEADER),
                HeaderValue::from_str(endpoint_signature)
                    .expect("endpoint signature should be a valid header value"),
            );
        }
    }
}

pub(super) fn finalize_gateway_response_with_context(
    state: &AppState,
    response: Response<Body>,
    remote_addr: &std::net::SocketAddr,
    request_context: &GatewayPublicRequestContext,
    execution_path: &'static str,
    started_at: &Instant,
    request_permit: Option<AdmissionPermit>,
) -> Response<Body> {
    finalize_gateway_response(
        state,
        response,
        &request_context.trace_id,
        remote_addr,
        &request_context.request_method,
        &request_context.request_path_and_query(),
        request_context.control_decision.as_ref(),
        execution_path,
        started_at,
        request_permit,
    )
}