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

systemprompt_api/services/gateway/service/
mod.rs
