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