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

systemprompt_api/services/gateway/service/
mod.rs
