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
14pub(super) use self::finalize::run_response_safety_scan;
15
16#[cfg(feature = "test-api")]
17pub mod test_api {
18    pub use super::blocks_at_phase;
19    pub use super::finalize::{apply_system_prompt_override, attach_request_id, dedupe_findings};
20}
21
22use std::sync::Arc;
23
24use anyhow::{Result, anyhow};
25use axum::body::Body;
26use axum::response::Response;
27use bytes::Bytes;
28use systemprompt_ai::{PHASE_REQUEST, PHASE_REQUEST_HISTORY, SafetyConfig, SafetyHistoryMode};
29use systemprompt_database::DbPool;
30use systemprompt_identifiers::{AiRequestId, UserId};
31use systemprompt_models::profile::{GatewayConfig, ProviderRegistry};
32
33use self::finalize::{
34    FinalizeCtx, apply_system_prompt_override, attach_request_id, finalize, run_request_safety_scan,
35};
36use self::resolve::{ResolvedUpstream, resolve_upstream};
37use super::audit::{GatewayAudit, GatewayRequestContext};
38use super::policy::{PolicyResolver, QuotaWindow};
39use super::protocol::canonical::CanonicalRequest;
40use super::protocol::inbound::InboundAdapter;
41use super::protocol::outbound::{OutboundCtx, OutboundOutcome, PreparedBody};
42use super::quota;
43use systemprompt_identifiers::{CallId, SessionId};
44use systemprompt_models::services::ai::ModelLimits;
45use systemprompt_models::wire::inspect;
46use systemprompt_security::authz::types::Decision;
47use systemprompt_security::policy::{
48    AgentScope, AuditOrigin, AuditTarget, ChainEntryResult, DecisionAudit, Evaluation,
49    GovernanceEngine, GovernedInput, GovernedTarget, PolicyContext, PrincipalSnapshot,
50    record_decision,
51};
52
53pub const REQUEST_ID_HEADER: &str = "x-systemprompt-request-id";
54
55#[derive(Debug, Clone, Copy)]
56pub struct GatewayService;
57
58#[derive(Debug)]
59pub struct DispatchInputs {
60    pub request: CanonicalRequest,
61    pub raw_body: Bytes,
62    pub ctx: GatewayRequestContext,
63    pub inbound: Arc<dyn InboundAdapter>,
64    pub forward_headers: Vec<(String, String)>,
65    pub identity_headers: Vec<(String, String)>,
66}
67
68#[derive(Debug, thiserror::Error)]
69pub enum DispatchError {
70    #[error(transparent)]
71    PreAudit(anyhow::Error),
72    #[error(transparent)]
73    Recorded(anyhow::Error),
74}
75
76#[derive(Debug, thiserror::Error)]
77#[error("{0}")]
78pub struct PolicyDenied(pub String);
79
80#[derive(Debug, thiserror::Error)]
81#[error("{message}")]
82pub struct QuotaExceeded {
83    pub message: String,
84    pub retry_after_seconds: i32,
85}
86
87#[derive(Debug, thiserror::Error)]
88#[error("{message}")]
89pub struct GuardForbidden {
90    pub message: String,
91}
92
93/// A denial from the typed four-stage governance chain — the same engine and
94/// the same operator-configured policies that govern MCP tool calls.
95#[derive(Debug, thiserror::Error)]
96#[error("{message}")]
97pub struct GovernanceDenied {
98    pub policy: String,
99    pub message: String,
100}
101
102#[derive(Debug, thiserror::Error)]
103#[error("{message}")]
104pub struct SafetyBlocked {
105    pub category: String,
106    pub message: String,
107}
108
109impl GatewayService {
110    pub async fn dispatch(
111        config: &GatewayConfig,
112        registry: &ProviderRegistry,
113        db: &DbPool,
114        repos: &super::GatewayRepositories,
115        inputs: DispatchInputs,
116    ) -> Result<Response<Body>, DispatchError> {
117        let DispatchInputs {
118            mut request,
119            raw_body,
120            ctx,
121            inbound,
122            forward_headers,
123            identity_headers,
124        } = inputs;
125        if ctx.session_id.is_none() {
126            return Err(DispatchError::PreAudit(anyhow!(
127                "gateway dispatch missing conversation binding (session_id)"
128            )));
129        }
130
131        let ai_request_id = ctx.ai_request_id.clone();
132        let upstream = resolve_upstream(config, registry, &request, &ai_request_id).await?;
133
134        tracing::info!(
135            ai_request_id = %ai_request_id,
136            user_id = %ctx.user_id,
137            model = %request.model,
138            provider = %upstream.route.provider,
139            upstream = %upstream.provider.endpoint,
140            wire_protocol = %ctx.wire_protocol,
141            streaming = request.stream,
142            "Gateway request dispatched"
143        );
144
145        let resolver = PolicyResolver::from_repository(repos.gateway_policies.clone());
146        let policy = resolver.resolve().await;
147
148        let audit = open_audit(repos, &ctx, &request, &raw_body, &identity_headers).await?;
149
150        if let Some(descriptor) = upstream.route_match_descriptor.as_deref() {
151            audit.set_route_match(descriptor).await;
152        }
153
154        enforce_quota(db, repos, &ctx.user_id, &policy.quota_windows, &audit).await?;
155        enforce_request_guards(db, &ctx.user_id, &upstream, &request, &audit).await?;
156
157        // Why: the payload is built before the scan so governance inspects the
158        // exact bytes that will go on the wire. Scanning the canonical form and
159        // sending something derived from it separately is how the two drift.
160        let prepared = prepare_payload(
161            config,
162            &upstream,
163            &mut request,
164            &audit,
165            UpstreamRelay {
166                raw_body: &raw_body,
167                inbound: inbound.as_ref(),
168            },
169        )
170        .await?;
171        audit.set_prepared_body_digest(&prepared.body.bytes).await;
172        attach_forwarded_surface(&mut request, &prepared, &ai_request_id);
173
174        // Why: governance runs before the scanner plane so first-deny-wins holds
175        // across both — a denied request never reaches the scanners, and so
176        // produces exactly one audit row and one 403.
177        enforce_governance(db, &ctx, &request, &audit).await?;
178        enforce_request_safety(repos, &ai_request_id, &request, &policy.safety, &audit).await?;
179
180        let outcome =
181            send_to_upstream(&upstream, &request, &prepared, &forward_headers, &audit).await?;
182
183        let response = finalize(
184            outcome,
185            FinalizeCtx {
186                audit: Arc::clone(&audit),
187                db: db.clone(),
188                repos: repos.clone(),
189                ai_request_id: ai_request_id.clone(),
190                policy,
191                inbound,
192                request_model: request.model.clone(),
193            },
194        )
195        .await;
196        Ok(attach_request_id(response, &ai_request_id))
197    }
198}
199
200struct UpstreamRelay<'a> {
201    raw_body: &'a Bytes,
202    inbound: &'a dyn InboundAdapter,
203}
204
205async fn open_audit(
206    repos: &super::GatewayRepositories,
207    ctx: &GatewayRequestContext,
208    request: &CanonicalRequest,
209    raw_body: &Bytes,
210    identity_headers: &[(String, String)],
211) -> Result<Arc<GatewayAudit>, DispatchError> {
212    let audit = Arc::new(GatewayAudit::new(repos, ctx.clone()));
213    if let Err(e) = audit.open(request, raw_body).await {
214        tracing::error!(error = %e, "audit open failed — proceeding without audit row");
215    }
216    // Why: these headers identify the client, user, and any spawning agent.
217    // They are recorded here, against the audit row's request id, and then
218    // dropped before the upstream send so a third-party provider never receives
219    // them. Emitted on the trace rather than as an `ai_requests` column because
220    // the attribution they add is per-agent, not per-request.
221    if !identity_headers.is_empty() {
222        tracing::info!(
223            ai_request_id = %ctx.ai_request_id,
224            user_id = %ctx.user_id,
225            headers = ?identity_headers,
226            "Gateway consumed client identity headers"
227        );
228    }
229    Ok(audit)
230}
231
232#[derive(Clone, Copy)]
233struct CtxParts<'a> {
234    upstream_model: &'a str,
235    model_limits: Option<ModelLimits>,
236    forward_headers: &'a [(String, String)],
237    raw_body: Option<&'a Bytes>,
238}
239
240struct Prepared {
241    upstream_model: String,
242    model_limits: Option<ModelLimits>,
243    body: PreparedBody,
244}
245
246async fn prepare_payload(
247    config: &GatewayConfig,
248    upstream: &ResolvedUpstream<'_>,
249    request: &mut CanonicalRequest,
250    audit: &GatewayAudit,
251    relay: UpstreamRelay<'_>,
252) -> Result<Prepared, DispatchError> {
253    let upstream_model = upstream
254        .route
255        .effective_upstream_model(&request.model)
256        .to_owned();
257    let override_descriptor =
258        apply_system_prompt_override(config, &upstream.provider.name, &upstream_model, request)
259            .await;
260    if let Some(descriptor) = &override_descriptor {
261        audit.set_system_prompt_override(descriptor).await;
262    }
263    let model_limits = upstream
264        .provider
265        .find_model(&upstream_model)
266        .map(|m| m.limits);
267    // Why: an applied override rewrote the canonical request, so the caller's
268    // original bytes no longer describe what the gateway decided to send.
269    let raw_body = (override_descriptor.is_none()
270        && relay.inbound.passthrough_wire() == Some(upstream.provider.wire))
271    .then_some(relay.raw_body);
272    // Why: unconditional, because an adapter may decline the raw lane and fall
273    // back to the canonical build. Stripping only when the lane was rejected up
274    // front would leave that fallback forwarding the identity.
275    strip_caller_identity(request);
276
277    let ctx = outbound_ctx(
278        upstream,
279        request,
280        CtxParts {
281            upstream_model: &upstream_model,
282            model_limits,
283            forward_headers: &[],
284            raw_body,
285        },
286    );
287    let body = upstream
288        .adapter
289        .build_body(&ctx)
290        .map_err(DispatchError::Recorded)?;
291    Ok(Prepared {
292        upstream_model,
293        model_limits,
294        body,
295    })
296}
297
298fn attach_forwarded_surface(
299    request: &mut CanonicalRequest,
300    prepared: &Prepared,
301    ai_request_id: &AiRequestId,
302) {
303    let surface = inspect::string_leaves(&prepared.body.bytes, inspect::SurfaceBudget::default());
304    if surface.truncated() {
305        tracing::warn!(
306            ai_request_id = %ai_request_id,
307            leaves = surface.len(),
308            "Gateway inspection surface truncated — part of the forwarded body was not scanned"
309        );
310    }
311    request.forwarded_surface = surface;
312}
313
314fn outbound_ctx<'a>(
315    upstream: &'a ResolvedUpstream<'a>,
316    request: &'a CanonicalRequest,
317    parts: CtxParts<'a>,
318) -> OutboundCtx<'a> {
319    OutboundCtx {
320        route: upstream.route.as_ref(),
321        endpoint: &upstream.provider.endpoint,
322        api_key: upstream.api_key,
323        request,
324        upstream_model: parts.upstream_model,
325        model_limits: parts.model_limits,
326        forward_headers: parts.forward_headers,
327        raw_body: parts.raw_body,
328    }
329}
330
331async fn send_to_upstream(
332    upstream: &ResolvedUpstream<'_>,
333    request: &CanonicalRequest,
334    prepared: &Prepared,
335    forward_headers: &[(String, String)],
336    audit: &GatewayAudit,
337) -> Result<OutboundOutcome, DispatchError> {
338    let ctx = outbound_ctx(
339        upstream,
340        request,
341        CtxParts {
342            upstream_model: &prepared.upstream_model,
343            model_limits: prepared.model_limits,
344            forward_headers,
345            raw_body: None,
346        },
347    );
348    match upstream.adapter.send(ctx, &prepared.body).await {
349        Ok(o) => Ok(o),
350        Err(e) => {
351            audit_upstream_failure(audit, upstream.provider.name.as_str(), &request.model, &e)
352                .await;
353            Err(DispatchError::Recorded(e))
354        },
355    }
356}
357
358// Why: `metadata.user_id` is an end-user identifier meant for the provider the
359// caller chose, so it must not reach a different wire's upstream. The
360// passthrough lane applies the same rule to the raw body in
361// `normalize_raw_body`.
362fn strip_caller_identity(request: &mut CanonicalRequest) {
363    let Some(metadata) = request.metadata.as_mut() else {
364        return;
365    };
366    let Some(obj) = metadata.as_object_mut() else {
367        return;
368    };
369    obj.remove("user_id");
370    if obj.is_empty() {
371        request.metadata = None;
372    }
373}
374
375async fn enforce_quota(
376    db: &DbPool,
377    repos: &super::GatewayRepositories,
378    user_id: &UserId,
379    quota_windows: &[QuotaWindow],
380    audit: &GatewayAudit,
381) -> Result<(), DispatchError> {
382    let reservation = quota::precheck_and_reserve(db, &repos.quota_buckets, user_id, quota_windows)
383        .await
384        .map_err(DispatchError::Recorded)?;
385    let Some(decision) = reservation else {
386        return Ok(());
387    };
388    if decision.allow {
389        return Ok(());
390    }
391    let msg = decision.message;
392    if let Err(e) = audit.fail(&msg).await {
393        tracing::warn!(error = %e, "quota audit fail failed");
394    }
395    Err(DispatchError::Recorded(
396        QuotaExceeded {
397            message: msg,
398            retry_after_seconds: decision.window_seconds,
399        }
400        .into(),
401    ))
402}
403
404async fn enforce_request_guards(
405    db: &DbPool,
406    user_id: &UserId,
407    upstream: &ResolvedUpstream<'_>,
408    request: &CanonicalRequest,
409    audit: &GatewayAudit,
410) -> Result<(), DispatchError> {
411    let Some(pool) = db.pool() else {
412        return Ok(());
413    };
414    let guard_request = systemprompt_extension::GatewayGuardRequest {
415        user_id: user_id.as_str(),
416        model: &request.model,
417        route_id: Some(upstream.route.id.as_str()),
418        provider: upstream.route.provider.as_str(),
419        streaming: request.stream,
420    };
421    let Err(deny) = systemprompt_extension::run_gateway_guards(&pool, &guard_request).await else {
422        return Ok(());
423    };
424    tracing::warn!(
425        user_id = %user_id,
426        model = %request.model,
427        route_id = %upstream.route.id,
428        kind = ?deny.kind,
429        reason = %deny.message,
430        "Gateway request denied by request guard"
431    );
432    if let Err(e) = audit.fail(&deny.message).await {
433        tracing::warn!(error = %e, "request-guard audit fail failed");
434    }
435    let inner: anyhow::Error = match deny.kind {
436        systemprompt_extension::GatewayDenyKind::Forbidden => GuardForbidden {
437            message: deny.message,
438        }
439        .into(),
440        systemprompt_extension::GatewayDenyKind::Quota => QuotaExceeded {
441            message: deny.message,
442            retry_after_seconds: deny.retry_after_seconds,
443        }
444        .into(),
445    };
446    Err(DispatchError::Recorded(inner))
447}
448
449// Why: written on allow as well as on deny — the chain trace is the product,
450// not just the refusals.
451async fn record_governance_decision(
452    db: &DbPool,
453    ctx: &GatewayRequestContext,
454    evaluation: Evaluation,
455    call_id: CallId,
456    session_id: SessionId,
457) {
458    let decision_audit = DecisionAudit {
459        id: uuid::Uuid::new_v4().to_string(),
460        call_id: call_id.as_str().to_owned(),
461        origin: AuditOrigin::Governed,
462        decision: evaluation.decision,
463        principal: PrincipalSnapshot {
464            user_id: ctx.user_id.clone(),
465            session_id,
466            agent_session: None,
467            agent_id: None,
468            agent_scope: ctx.access_scope,
469        },
470        target: AuditTarget {
471            tool_name: GovernedTarget::Prompt.as_str().to_owned(),
472            plugin_id: None,
473        },
474        chain: evaluation.chain,
475        approver: None,
476        act_chain: Vec::new(),
477        context_id: Some(ctx.context_id.as_str().to_owned()),
478        trace_id: ctx.trace_id.as_ref().map(|t| t.as_str().to_owned()),
479    };
480    match db.write_pool_arc() {
481        Ok(pool) => {
482            if let Err(e) = record_decision(&pool, &decision_audit).await {
483                tracing::error!(
484                    target: "governance.audit.write_failed",
485                    error = %e,
486                    ai_request_id = %ctx.ai_request_id,
487                    "gateway governance audit write failed; row dropped"
488                );
489            }
490        },
491        Err(e) => tracing::error!(
492            target: "governance.audit.write_failed",
493            error = %e,
494            ai_request_id = %ctx.ai_request_id,
495            "no write pool for the gateway governance decision; row dropped"
496        ),
497    }
498}
499
500struct PromptEvaluation {
501    evaluation: Evaluation,
502    call_id: CallId,
503    session_id: SessionId,
504}
505
506fn evaluate_prompt(ctx: &GatewayRequestContext, request: &CanonicalRequest) -> PromptEvaluation {
507    // Why: `flatten_text` includes the forwarded surface attached just above,
508    // so the chain scans exactly the bytes that will go on the wire — operator
509    // `extra_patterns` included, which the hardcoded safety scanner cannot do.
510    let input = GovernedInput::prompt(request.flatten_text());
511    // Why: the bucket key is `session_id:user_id`. A sessionless inference call
512    // needs a *stable* placeholder — minting one per request would give every
513    // call its own bucket and silently disable rate limiting.
514    let session_id = ctx.session_id.clone().unwrap_or_else(SessionId::system);
515    // Why: the engine's idempotency contract is per-call_id, and the ai-request
516    // id is the one identifier stable across re-evaluations of this call. It is
517    // what stops the rate limiter charging twice.
518    let call_id = CallId::new(ctx.ai_request_id.as_str());
519
520    // Why: the same engine instance the MCP governance webhook uses — the rate
521    // limiter's buckets are instance-scoped, so a second engine would give
522    // inference its own budget and silently double every operator limit.
523    let evaluation = GovernanceEngine::global().evaluate(&PolicyContext {
524        target: GovernedTarget::Prompt,
525        agent_scope: AgentScope::User {
526            user_id: ctx.user_id.clone(),
527        },
528        access_scope: ctx.access_scope,
529        session_id: &session_id,
530        user_id: &ctx.user_id,
531        input: &input,
532        call_id: &call_id,
533    });
534
535    PromptEvaluation {
536        evaluation,
537        call_id,
538        session_id,
539    }
540}
541
542async fn enforce_governance(
543    db: &DbPool,
544    ctx: &GatewayRequestContext,
545    request: &CanonicalRequest,
546    audit: &GatewayAudit,
547) -> Result<(), DispatchError> {
548    let PromptEvaluation {
549        evaluation,
550        call_id,
551        session_id,
552    } = evaluate_prompt(ctx, request);
553
554    let denied = match &evaluation.decision {
555        Decision::Allow { .. } => None,
556        Decision::Deny { reason } => Some(reason.to_string()),
557    };
558    let policy = evaluation
559        .chain
560        .iter()
561        .find(|e| e.result == ChainEntryResult::Fail)
562        .map_or_else(
563            || "default_allow".to_owned(),
564            |e| e.policy_id.as_str().to_owned(),
565        );
566
567    record_governance_decision(db, ctx, evaluation, call_id, session_id).await;
568
569    let Some(reason) = denied else {
570        return Ok(());
571    };
572    tracing::warn!(
573        ai_request_id = %ctx.ai_request_id,
574        user_id = %ctx.user_id,
575        policy = %policy,
576        reason = %reason,
577        "Gateway request denied by governance policy"
578    );
579    if let Err(e) = audit.fail(&reason).await {
580        tracing::warn!(error = %e, "governance-deny audit fail failed");
581    }
582    Err(DispatchError::Recorded(
583        GovernanceDenied {
584            policy,
585            message: reason,
586        }
587        .into(),
588    ))
589}
590
591async fn enforce_request_safety(
592    repos: &super::GatewayRepositories,
593    ai_request_id: &AiRequestId,
594    request: &CanonicalRequest,
595    safety: &SafetyConfig,
596    audit: &GatewayAudit,
597) -> Result<(), DispatchError> {
598    let findings =
599        run_request_safety_scan(&repos.safety_findings, ai_request_id, request, safety).await;
600    let Some(finding) = findings.iter().find(|f| {
601        safety.block_categories.contains(&f.category) && blocks_at_phase(f.phase, safety.history)
602    }) else {
603        return Ok(());
604    };
605    let msg = format!(
606        "request blocked by safety policy: category '{}'",
607        finding.category
608    );
609    tracing::warn!(
610        ai_request_id = %ai_request_id,
611        category = %finding.category,
612        scanner = %finding.scanner,
613        "Gateway blocked request by safety policy"
614    );
615    if let Err(e) = audit.fail(&msg).await {
616        tracing::warn!(error = %e, "safety-block audit fail failed");
617    }
618    Err(DispatchError::Recorded(
619        SafetyBlocked {
620            category: finding.category.clone(),
621            message: msg,
622        }
623        .into(),
624    ))
625}
626
627pub fn blocks_at_phase(phase: &str, history: SafetyHistoryMode) -> bool {
628    match phase {
629        PHASE_REQUEST => true,
630        PHASE_REQUEST_HISTORY => history == SafetyHistoryMode::Block,
631        _ => false,
632    }
633}
634
635async fn audit_upstream_failure(
636    audit: &GatewayAudit,
637    provider: &str,
638    model: &str,
639    error: &anyhow::Error,
640) {
641    tracing::warn!(
642        provider = %provider,
643        model = %model,
644        error = %error,
645        "gateway upstream call failed"
646    );
647    if let Err(audit_err) = audit.fail(&error.to_string()).await {
648        tracing::warn!(error = %audit_err, "upstream audit fail failed");
649    }
650}
systemprompt_api/services/gateway/service/mod.rs

systemprompt_api/services/gateway/service/
mod.rs
