Skip to main content

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;
13mod stages;
14
15pub(super) use self::finalize::run_response_safety_scan;
16
17#[cfg(feature = "test-api")]
18pub mod test_api {
19    pub use super::blocks_at_phase;
20    pub use super::finalize::{apply_system_prompt_override, attach_request_id, dedupe_findings};
21    pub use super::resolve::{describe_route_match, enforce_route_requirements};
22}
23
24use std::sync::Arc;
25
26use anyhow::{Result, anyhow};
27use axum::body::Body;
28use axum::response::Response;
29use bytes::Bytes;
30use systemprompt_ai::{PHASE_REQUEST, PHASE_REQUEST_HISTORY, SafetyHistoryMode};
31use systemprompt_database::DbPool;
32use systemprompt_identifiers::UserId;
33use systemprompt_models::profile::{GatewayConfig, ProviderRegistry};
34
35use self::finalize::{FinalizeCtx, attach_request_id, finalize};
36use self::resolve::{ResolvedUpstream, resolve_upstream};
37use self::stages::{GovernedDispatch, PreparedDispatch, ScannedDispatch, UpstreamRelay};
38use super::audit::{GatewayAudit, GatewayRequestContext};
39use super::policy::{PolicyResolver, QuotaWindow};
40use super::protocol::canonical::CanonicalRequest;
41use super::protocol::inbound::InboundAdapter;
42use super::quota;
43
44pub const REQUEST_ID_HEADER: &str = "x-systemprompt-request-id";
45
46#[derive(Debug, Clone, Copy)]
47pub struct GatewayService;
48
49#[derive(Debug)]
50pub struct DispatchInputs {
51    pub request: CanonicalRequest,
52    pub raw_body: Bytes,
53    pub ctx: GatewayRequestContext,
54    pub inbound: Arc<dyn InboundAdapter>,
55    pub forward_headers: Vec<(String, String)>,
56    pub identity_headers: Vec<(String, String)>,
57}
58
59#[derive(Debug, thiserror::Error)]
60pub enum DispatchError {
61    #[error(transparent)]
62    PreAudit(anyhow::Error),
63    #[error(transparent)]
64    Recorded(anyhow::Error),
65}
66
67#[derive(Debug, thiserror::Error)]
68#[error("{0}")]
69pub struct PolicyDenied(pub String);
70
71#[derive(Debug, thiserror::Error)]
72#[error("{message}")]
73pub struct QuotaExceeded {
74    pub message: String,
75    pub retry_after_seconds: i32,
76}
77
78#[derive(Debug, thiserror::Error)]
79#[error("{message}")]
80pub struct GuardForbidden {
81    pub message: String,
82}
83
84/// A denial from the typed four-stage governance chain — the same engine and
85/// the same operator-configured policies that govern MCP tool calls.
86#[derive(Debug, thiserror::Error)]
87#[error("{message}")]
88pub struct GovernanceDenied {
89    pub policy: String,
90    pub message: String,
91}
92
93#[derive(Debug, thiserror::Error)]
94#[error("{message}")]
95pub struct SafetyBlocked {
96    pub category: String,
97    pub message: String,
98}
99
100impl GatewayService {
101    pub async fn dispatch(
102        config: &GatewayConfig,
103        registry: &ProviderRegistry,
104        db: &DbPool,
105        repos: &super::GatewayRepositories,
106        inputs: DispatchInputs,
107    ) -> Result<Response<Body>, DispatchError> {
108        let DispatchInputs {
109            request,
110            raw_body,
111            ctx,
112            inbound,
113            forward_headers,
114            identity_headers,
115        } = inputs;
116        if ctx.session_id.is_none() {
117            return Err(DispatchError::PreAudit(anyhow!(
118                "gateway dispatch missing conversation binding (session_id)"
119            )));
120        }
121
122        let ai_request_id = ctx.ai_request_id.clone();
123        let upstream = resolve_upstream(config, registry, &request, &ai_request_id).await?;
124
125        tracing::info!(
126            ai_request_id = %ai_request_id,
127            user_id = %ctx.user_id,
128            model = %request.model,
129            provider = %upstream.route.provider,
130            upstream = %upstream.provider.endpoint,
131            wire_protocol = %ctx.wire_protocol,
132            streaming = request.stream,
133            "Gateway request dispatched"
134        );
135
136        let resolver = PolicyResolver::from_repository(repos.gateway_policies.clone());
137        let policy = resolver.resolve().await;
138
139        let audit = open_audit(repos, &ctx, &request, &raw_body, &identity_headers).await?;
140
141        if let Some(descriptor) = upstream.route_match_descriptor.as_deref() {
142            audit.set_route_match(descriptor).await;
143        }
144
145        enforce_quota(db, repos, &ctx.user_id, &policy.quota_windows, &audit).await?;
146        enforce_request_guards(db, &ctx.user_id, &upstream, &request, &audit).await?;
147
148        let prepared = PreparedDispatch::build(
149            config,
150            &upstream,
151            request,
152            &audit,
153            UpstreamRelay {
154                raw_body: &raw_body,
155                inbound: inbound.as_ref(),
156            },
157        )
158        .await?;
159        let governed = GovernedDispatch::enforce(prepared, db, &ctx, &audit).await?;
160        let scanned =
161            ScannedDispatch::enforce(governed, repos, &ai_request_id, &policy.safety, &audit)
162                .await?;
163
164        let outcome = scanned.send(&upstream, &forward_headers, &audit).await?;
165
166        let response = finalize(
167            outcome,
168            FinalizeCtx {
169                audit: Arc::clone(&audit),
170                db: db.clone(),
171                repos: repos.clone(),
172                ai_request_id: ai_request_id.clone(),
173                policy,
174                inbound,
175                request_model: scanned.request_model().to_owned(),
176            },
177        )
178        .await;
179        Ok(attach_request_id(response, &ai_request_id))
180    }
181}
182
183async fn open_audit(
184    repos: &super::GatewayRepositories,
185    ctx: &GatewayRequestContext,
186    request: &CanonicalRequest,
187    raw_body: &Bytes,
188    identity_headers: &[(String, String)],
189) -> Result<Arc<GatewayAudit>, DispatchError> {
190    let audit = Arc::new(GatewayAudit::new(repos, ctx.clone()));
191    if let Err(e) = audit.open(request, raw_body).await {
192        tracing::error!(error = %e, "audit open failed — proceeding without audit row");
193    }
194    // Why: identity headers are recorded against the audit row, then dropped
195    // before the upstream send so a third-party provider never receives them.
196    if !identity_headers.is_empty() {
197        tracing::info!(
198            ai_request_id = %ctx.ai_request_id,
199            user_id = %ctx.user_id,
200            headers = ?identity_headers,
201            "Gateway consumed client identity headers"
202        );
203    }
204    Ok(audit)
205}
206
207async fn enforce_quota(
208    db: &DbPool,
209    repos: &super::GatewayRepositories,
210    user_id: &UserId,
211    quota_windows: &[QuotaWindow],
212    audit: &GatewayAudit,
213) -> Result<(), DispatchError> {
214    let reservation = quota::precheck_and_reserve(db, &repos.quota_buckets, user_id, quota_windows)
215        .await
216        .map_err(DispatchError::Recorded)?;
217    let Some(decision) = reservation else {
218        return Ok(());
219    };
220    if decision.allow {
221        return Ok(());
222    }
223    let msg = decision.message;
224    if let Err(e) = audit.fail(&msg).await {
225        tracing::warn!(error = %e, "quota audit fail failed");
226    }
227    Err(DispatchError::Recorded(
228        QuotaExceeded {
229            message: msg,
230            retry_after_seconds: decision.window_seconds,
231        }
232        .into(),
233    ))
234}
235
236async fn enforce_request_guards(
237    db: &DbPool,
238    user_id: &UserId,
239    upstream: &ResolvedUpstream<'_>,
240    request: &CanonicalRequest,
241    audit: &GatewayAudit,
242) -> Result<(), DispatchError> {
243    let Some(pool) = db.pool() else {
244        return Ok(());
245    };
246    let guard_request = systemprompt_extension::GatewayGuardRequest {
247        user_id: user_id.as_str(),
248        model: &request.model,
249        route_id: Some(upstream.route.id.as_str()),
250        provider: upstream.route.provider.as_str(),
251        streaming: request.stream,
252    };
253    let Err(deny) = systemprompt_extension::run_gateway_guards(&pool, &guard_request).await else {
254        return Ok(());
255    };
256    tracing::warn!(
257        user_id = %user_id,
258        model = %request.model,
259        route_id = %upstream.route.id,
260        kind = ?deny.kind,
261        reason = %deny.message,
262        "Gateway request denied by request guard"
263    );
264    if let Err(e) = audit.fail(&deny.message).await {
265        tracing::warn!(error = %e, "request-guard audit fail failed");
266    }
267    let inner: anyhow::Error = match deny.kind {
268        systemprompt_extension::GatewayDenyKind::Forbidden => GuardForbidden {
269            message: deny.message,
270        }
271        .into(),
272        systemprompt_extension::GatewayDenyKind::Quota => QuotaExceeded {
273            message: deny.message,
274            retry_after_seconds: deny.retry_after_seconds,
275        }
276        .into(),
277    };
278    Err(DispatchError::Recorded(inner))
279}
280
281pub fn blocks_at_phase(phase: &str, history: SafetyHistoryMode) -> bool {
282    match phase {
283        PHASE_REQUEST => true,
284        PHASE_REQUEST_HISTORY => history == SafetyHistoryMode::Block,
285        _ => false,
286    }
287}