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