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