systemprompt-api 0.51.0

//! Gateway message-dispatch entry point.
//!
//! [`handle`] is the shared handler behind the `/messages` and `/responses`
//! routes: it builds a `RequestContext`, extracts and authorizes the request
//! (`extract`), then dispatches to the resolved provider (`dispatch`),
//! persisting a rejection record (`rejection`) on any early failure. Inbound
//! wire format is selected by the [`InboundAdapter`] passed in by the router.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

pub mod auth;
pub mod dispatch;
pub mod extract;
pub mod rejection;

pub use dispatch::map_upstream_error;
pub use extract::{GatewayAuthzRequestInput, build_gateway_authz_request, extract_credential};

use axum::body::Body;
use axum::extract::Request;
use axum::http::StatusCode;
use axum::response::Response;
use std::sync::Arc;
use systemprompt_identifiers::AiRequestId;
use systemprompt_loader::ServicesBootstrap;
use systemprompt_models::services::ServicesConfig;
use systemprompt_runtime::AppContext;

use crate::services::gateway::audit::GatewayAccessLog;
use crate::services::gateway::protocol::inbound::InboundAdapter;
use crate::services::middleware::JwtContextExtractor;

use dispatch::{RejectionError, build_error_response, dispatch_to_provider, error_type_for};
use extract::{RejectionPartial, extract_request_context};
use rejection::persist_rejection;

pub(super) struct RequestContext<'a> {
    pub jwt_extractor: &'a JwtContextExtractor,
    pub ctx: &'a AppContext,
    pub repos: &'a crate::services::gateway::GatewayRepositories,
    pub services: &'static ServicesConfig,
    pub ai_request_id: &'a AiRequestId,
    pub access_log: Option<GatewayAccessLog>,
}

pub async fn handle(
    inbound: Arc<dyn InboundAdapter>,
    jwt_extractor: Arc<JwtContextExtractor>,
    ctx: AppContext,
    repos: Arc<crate::services::gateway::GatewayRepositories>,
    request: Request<Body>,
) -> Response<Body> {
    let ai_request_id = AiRequestId::generate();
    let mut partial = RejectionPartial::default();
    let inner = HandleInner {
        inbound: Arc::clone(&inbound),
        jwt_extractor: &jwt_extractor,
        ctx: &ctx,
        repos: &repos,
        ai_request_id: &ai_request_id,
        partial: &mut partial,
    };
    let mut response = match inner.run(request).await {
        Ok(resp) => resp,
        Err(RejectionError {
            status,
            message,
            persist,
        }) => {
            tracing::warn!(
                status = %status,
                message = %message,
                ai_request_id = %ai_request_id,
                wire = inbound.wire_name(),
                "Gateway request rejected",
            );
            if persist {
                persist_rejection(&repos, &ai_request_id, &partial, status, &message).await;
            }
            let body = inbound.render_error(status, &message);
            Response::builder()
                .status(status)
                .header("content-type", "application/json")
                .body(Body::from(body))
                .unwrap_or_else(|_| build_error_response(status, error_type_for(status), &message))
        },
    };
    attach_log_identity(&mut response, &partial);
    response
}

fn attach_log_identity(response: &mut Response<Body>, partial: &RejectionPartial) {
    let (Some(user_id), Some(session_id), Some(trace_id)) = (
        partial.user_id.as_ref(),
        partial.session_id.as_ref(),
        partial.trace_id.as_ref(),
    ) else {
        return;
    };
    response.extensions_mut().insert(super::GatewayLogIdentity {
        user: user_id.clone(),
        session: session_id.clone(),
        trace: trace_id.clone(),
    });
}

struct HandleInner<'a> {
    inbound: Arc<dyn InboundAdapter>,
    jwt_extractor: &'a JwtContextExtractor,
    ctx: &'a AppContext,
    repos: &'a crate::services::gateway::GatewayRepositories,
    ai_request_id: &'a AiRequestId,
    partial: &'a mut RejectionPartial,
}

impl HandleInner<'_> {
    async fn run(self, request: Request<Body>) -> Result<Response<Body>, RejectionError> {
        let services = ServicesBootstrap::get().map_err(|e| RejectionError {
            status: StatusCode::SERVICE_UNAVAILABLE,
            message: format!("Services config not ready: {e}"),
            persist: true,
        })?;
        let access_log = request.extensions().get::<GatewayAccessLog>().cloned();
        let request_ctx = RequestContext {
            jwt_extractor: self.jwt_extractor,
            ctx: self.ctx,
            repos: self.repos,
            services,
            ai_request_id: self.ai_request_id,
            access_log,
        };
        let prepared = extract_request_context(&request_ctx, &self.inbound, request, self.partial)
            .await
            .map_err(|(status, message)| RejectionError {
                status,
                message,
                persist: true,
            })?;
        dispatch_to_provider(&request_ctx, self.inbound, prepared).await
    }
}