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

systemprompt_api/services/gateway/service/
mod.rs
