systemprompt_api/services/gateway/service/
mod.rs

1//! Gateway dispatch entry point: route resolution, policy and quota checks,
2//! upstream send, and response finalization.
3#![expect(
4    clippy::clone_on_ref_ptr,
5    reason = "Arc::clone usage is intentional and ergonomic in this gateway dispatch path"
6)]
7
8mod finalize;
9mod resolve;
10
11use std::sync::Arc;
12
13use anyhow::{Result, anyhow};
14use axum::body::Body;
15use axum::response::Response;
16use bytes::Bytes;
17use systemprompt_ai::SafetyConfig;
18use systemprompt_database::DbPool;
19use systemprompt_identifiers::{AiRequestId, UserId};
20use systemprompt_models::profile::{GatewayConfig, ProviderRegistry};
21
22use self::finalize::{
23    FinalizeCtx, apply_system_prompt_override, attach_request_id, finalize, run_request_safety_scan,
24};
25use self::resolve::{ResolvedUpstream, resolve_upstream};
26use super::audit::{GatewayAudit, GatewayRequestContext};
27use super::policy::{PolicyResolver, QuotaWindow};
28use super::protocol::canonical::CanonicalRequest;
29use super::protocol::inbound::InboundAdapter;
30use super::protocol::outbound::{OutboundCtx, OutboundOutcome};
31use super::quota;
32
33pub const REQUEST_ID_HEADER: &str = "x-systemprompt-request-id";
34
35#[derive(Debug, Clone, Copy)]
36pub struct GatewayService;
37
38#[derive(Debug)]
39pub struct DispatchInputs {
40    pub request: CanonicalRequest,
41    pub raw_body: Bytes,
42    pub ctx: GatewayRequestContext,
43    pub inbound: Arc<dyn InboundAdapter>,
44}
45
46#[derive(Debug, thiserror::Error)]
47pub enum DispatchError {
48    #[error(transparent)]
49    PreAudit(anyhow::Error),
50    #[error(transparent)]
51    Recorded(anyhow::Error),
52}
53
54#[derive(Debug, thiserror::Error)]
55#[error("{0}")]
56pub struct PolicyDenied(pub String);
57
58#[derive(Debug, thiserror::Error)]
59#[error("{message}")]
60pub struct QuotaExceeded {
61    pub message: String,
62    pub retry_after_seconds: i32,
63}
64
65#[derive(Debug, thiserror::Error)]
66#[error("{message}")]
67pub struct SafetyBlocked {
68    pub category: String,
69    pub message: String,
70}
71
72impl GatewayService {
73    pub async fn dispatch(
74        config: &GatewayConfig,
75        registry: &ProviderRegistry,
76        db: &DbPool,
77        inputs: DispatchInputs,
78    ) -> Result<Response<Body>, DispatchError> {
79        let DispatchInputs {
80            mut request,
81            raw_body,
82            ctx,
83            inbound,
84        } = inputs;
85        if ctx.session_id.is_none() {
86            return Err(DispatchError::PreAudit(anyhow!(
87                "gateway dispatch missing conversation binding (session_id)"
88            )));
89        }
90
91        let ai_request_id = ctx.ai_request_id.clone();
92        let upstream = resolve_upstream(config, registry, &request, &ai_request_id).await?;
93
94        tracing::info!(
95            ai_request_id = %ai_request_id,
96            user_id = %ctx.user_id,
97            model = %request.model,
98            provider = %upstream.route.provider,
99            upstream = %upstream.provider.endpoint,
100            wire_protocol = %ctx.wire_protocol,
101            streaming = request.stream,
102            "Gateway request dispatched"
103        );
104
105        let resolver = PolicyResolver::new(db).map_err(DispatchError::PreAudit)?;
106        let policy = resolver.resolve().await;
107
108        let audit = Arc::new(
109            GatewayAudit::new(db, ctx.clone())
110                .map_err(|e| DispatchError::PreAudit(anyhow!("audit init failed: {e}")))?,
111        );
112
113        if let Err(e) = audit.open(&request, &raw_body).await {
114            tracing::error!(error = %e, "audit open failed — proceeding without audit row");
115        }
116
117        if let Some(descriptor) = upstream.route_match_descriptor.as_deref() {
118            audit.set_route_match(descriptor).await;
119        }
120
121        enforce_quota(db, &ctx.user_id, &policy.quota_windows, &audit).await?;
122        enforce_request_safety(db, &ai_request_id, &request, &policy.safety, &audit).await?;
123
124        let outcome = send_to_upstream(config, &upstream, &mut request, &audit).await?;
125
126        let response = finalize(
127            outcome,
128            FinalizeCtx {
129                audit: Arc::clone(&audit),
130                db: db.clone(),
131                ai_request_id: ai_request_id.clone(),
132                policy,
133                inbound,
134                request_model: request.model.clone(),
135            },
136        )
137        .await;
138        Ok(attach_request_id(response, &ai_request_id))
139    }
140}
141
142async fn send_to_upstream(
143    config: &GatewayConfig,
144    upstream: &ResolvedUpstream<'_>,
145    request: &mut CanonicalRequest,
146    audit: &GatewayAudit,
147) -> Result<OutboundOutcome, DispatchError> {
148    let upstream_model = upstream
149        .route
150        .effective_upstream_model(&request.model)
151        .to_owned();
152    if let Some(descriptor) =
153        apply_system_prompt_override(config, &upstream.provider.name, &upstream_model, request)
154            .await
155    {
156        audit.set_system_prompt_override(&descriptor).await;
157    }
158    let model_limits = upstream
159        .provider
160        .find_model(&upstream_model)
161        .map(|m| m.limits);
162    let outbound_ctx = OutboundCtx {
163        route: upstream.route.as_ref(),
164        endpoint: &upstream.provider.endpoint,
165        api_key: upstream.api_key,
166        request,
167        upstream_model: &upstream_model,
168        model_limits,
169    };
170
171    match upstream.adapter.send(outbound_ctx).await {
172        Ok(o) => Ok(o),
173        Err(e) => {
174            audit_upstream_failure(audit, upstream.provider.name.as_str(), &request.model, &e)
175                .await;
176            Err(DispatchError::Recorded(e))
177        },
178    }
179}
180
181async fn enforce_quota(
182    db: &DbPool,
183    user_id: &UserId,
184    quota_windows: &[QuotaWindow],
185    audit: &GatewayAudit,
186) -> Result<(), DispatchError> {
187    let reservation = quota::precheck_and_reserve(db, user_id, quota_windows)
188        .await
189        .map_err(DispatchError::Recorded)?;
190    let Some(decision) = reservation else {
191        return Ok(());
192    };
193    if decision.allow {
194        return Ok(());
195    }
196    let msg = format!(
197        "quota exceeded for window {}s (used {}/{:?})",
198        decision.window_seconds, decision.state.requests, decision.limit_requests
199    );
200    if let Err(e) = audit.fail(&msg).await {
201        tracing::warn!(error = %e, "quota audit fail failed");
202    }
203    Err(DispatchError::Recorded(
204        QuotaExceeded {
205            message: msg,
206            retry_after_seconds: decision.window_seconds,
207        }
208        .into(),
209    ))
210}
211
212async fn enforce_request_safety(
213    db: &DbPool,
214    ai_request_id: &AiRequestId,
215    request: &CanonicalRequest,
216    safety: &SafetyConfig,
217    audit: &GatewayAudit,
218) -> Result<(), DispatchError> {
219    let findings = run_request_safety_scan(db, ai_request_id, request, safety).await;
220    let Some(finding) = findings
221        .iter()
222        .find(|f| safety.block_categories.contains(&f.category))
223    else {
224        return Ok(());
225    };
226    let msg = format!(
227        "request blocked by safety policy: category '{}'",
228        finding.category
229    );
230    tracing::warn!(
231        ai_request_id = %ai_request_id,
232        category = %finding.category,
233        scanner = %finding.scanner,
234        "Gateway blocked request by safety policy"
235    );
236    if let Err(e) = audit.fail(&msg).await {
237        tracing::warn!(error = %e, "safety-block audit fail failed");
238    }
239    Err(DispatchError::Recorded(
240        SafetyBlocked {
241            category: finding.category.clone(),
242            message: msg,
243        }
244        .into(),
245    ))
246}
247
248async fn audit_upstream_failure(
249    audit: &GatewayAudit,
250    provider: &str,
251    model: &str,
252    error: &anyhow::Error,
253) {
254    tracing::warn!(
255        provider = %provider,
256        model = %model,
257        error = %error,
258        "gateway upstream call failed"
259    );
260    if let Err(audit_err) = audit.fail(&error.to_string()).await {
261        tracing::warn!(error = %audit_err, "upstream audit fail failed");
262    }
263}
systemprompt_api/services/gateway/service/mod.rs

systemprompt_api/services/gateway/service/
mod.rs
