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

systemprompt_api/services/gateway/service/
mod.rs
