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        inputs: DispatchInputs,
119    ) -> Result<Response<Body>, DispatchError> {
120        let DispatchInputs {
121            mut request,
122            raw_body,
123            ctx,
124            inbound,
125            forward_headers,
126            identity_headers,
127        } = inputs;
128        if ctx.session_id.is_none() {
129            return Err(DispatchError::PreAudit(anyhow!(
130                "gateway dispatch missing conversation binding (session_id)"
131            )));
132        }
133
134        let ai_request_id = ctx.ai_request_id.clone();
135        let upstream = resolve_upstream(config, registry, &request, &ai_request_id).await?;
136
137        tracing::info!(
138            ai_request_id = %ai_request_id,
139            user_id = %ctx.user_id,
140            model = %request.model,
141            provider = %upstream.route.provider,
142            upstream = %upstream.provider.endpoint,
143            wire_protocol = %ctx.wire_protocol,
144            streaming = request.stream,
145            "Gateway request dispatched"
146        );
147
148        let resolver = PolicyResolver::new(db).map_err(DispatchError::PreAudit)?;
149        let policy = resolver.resolve().await;
150
151        let audit = open_audit(db, &ctx, &request, &raw_body, &identity_headers).await?;
152
153        if let Some(descriptor) = upstream.route_match_descriptor.as_deref() {
154            audit.set_route_match(descriptor).await;
155        }
156
157        enforce_quota(db, &ctx.user_id, &policy.quota_windows, &audit).await?;
158        enforce_request_guards(db, &ctx.user_id, &upstream, &request, &audit).await?;
159
160        // Why: the payload is built before the scan so governance inspects the
161        // exact bytes that will go on the wire. Scanning the canonical form and
162        // sending something derived from it separately is how the two drift.
163        let prepared = prepare_payload(
164            config,
165            &upstream,
166            &mut request,
167            &audit,
168            UpstreamRelay {
169                raw_body: &raw_body,
170                inbound: inbound.as_ref(),
171            },
172        )
173        .await?;
174        audit.set_prepared_body_digest(&prepared.body.bytes).await;
175        attach_forwarded_surface(&mut request, &prepared, &ai_request_id);
176
177        // Why: governance runs before the scanner plane so first-deny-wins holds
178        // across both — a denied request never reaches the scanners, and so
179        // produces exactly one audit row and one 403.
180        enforce_governance(db, &ctx, &request, &audit).await?;
181        enforce_request_safety(db, &ai_request_id, &request, &policy.safety, &audit).await?;
182
183        let outcome =
184            send_to_upstream(&upstream, &request, &prepared, &forward_headers, &audit).await?;
185
186        let response = finalize(
187            outcome,
188            FinalizeCtx {
189                audit: Arc::clone(&audit),
190                db: db.clone(),
191                ai_request_id: ai_request_id.clone(),
192                policy,
193                inbound,
194                request_model: request.model.clone(),
195            },
196        )
197        .await;
198        Ok(attach_request_id(response, &ai_request_id))
199    }
200}
201
202struct UpstreamRelay<'a> {
203    raw_body: &'a Bytes,
204    inbound: &'a dyn InboundAdapter,
205}
206
207async fn open_audit(
208    db: &DbPool,
209    ctx: &GatewayRequestContext,
210    request: &CanonicalRequest,
211    raw_body: &Bytes,
212    identity_headers: &[(String, String)],
213) -> Result<Arc<GatewayAudit>, DispatchError> {
214    let audit = Arc::new(
215        GatewayAudit::new(db, ctx.clone())
216            .map_err(|e| DispatchError::PreAudit(anyhow!("audit init failed: {e}")))?,
217    );
218    if let Err(e) = audit.open(request, raw_body).await {
219        tracing::error!(error = %e, "audit open failed — proceeding without audit row");
220    }
221    // Why: these headers identify the client, user, and any spawning agent.
222    // They are recorded here, against the audit row's request id, and then
223    // dropped before the upstream send so a third-party provider never receives
224    // them. Emitted on the trace rather than as an `ai_requests` column because
225    // the attribution they add is per-agent, not per-request.
226    if !identity_headers.is_empty() {
227        tracing::info!(
228            ai_request_id = %ctx.ai_request_id,
229            user_id = %ctx.user_id,
230            headers = ?identity_headers,
231            "Gateway consumed client identity headers"
232        );
233    }
234    Ok(audit)
235}
236
237#[derive(Clone, Copy)]
238struct CtxParts<'a> {
239    upstream_model: &'a str,
240    model_limits: Option<ModelLimits>,
241    forward_headers: &'a [(String, String)],
242    raw_body: Option<&'a Bytes>,
243}
244
245struct Prepared {
246    upstream_model: String,
247    model_limits: Option<ModelLimits>,
248    body: PreparedBody,
249}
250
251async fn prepare_payload(
252    config: &GatewayConfig,
253    upstream: &ResolvedUpstream<'_>,
254    request: &mut CanonicalRequest,
255    audit: &GatewayAudit,
256    relay: UpstreamRelay<'_>,
257) -> Result<Prepared, DispatchError> {
258    let upstream_model = upstream
259        .route
260        .effective_upstream_model(&request.model)
261        .to_owned();
262    let override_descriptor =
263        apply_system_prompt_override(config, &upstream.provider.name, &upstream_model, request)
264            .await;
265    if let Some(descriptor) = &override_descriptor {
266        audit.set_system_prompt_override(descriptor).await;
267    }
268    let model_limits = upstream
269        .provider
270        .find_model(&upstream_model)
271        .map(|m| m.limits);
272    // Why: an applied override rewrote the canonical request, so the caller's
273    // original bytes no longer describe what the gateway decided to send.
274    let raw_body = (override_descriptor.is_none()
275        && relay.inbound.passthrough_wire() == Some(upstream.provider.wire))
276    .then_some(relay.raw_body);
277    // Why: unconditional, because an adapter may decline the raw lane and fall
278    // back to the canonical build. Stripping only when the lane was rejected up
279    // front would leave that fallback forwarding the identity.
280    strip_caller_identity(request);
281
282    let ctx = outbound_ctx(
283        upstream,
284        request,
285        CtxParts {
286            upstream_model: &upstream_model,
287            model_limits,
288            forward_headers: &[],
289            raw_body,
290        },
291    );
292    let body = upstream
293        .adapter
294        .build_body(&ctx)
295        .map_err(DispatchError::Recorded)?;
296    Ok(Prepared {
297        upstream_model,
298        model_limits,
299        body,
300    })
301}
302
303fn attach_forwarded_surface(
304    request: &mut CanonicalRequest,
305    prepared: &Prepared,
306    ai_request_id: &AiRequestId,
307) {
308    let surface = inspect::string_leaves(&prepared.body.bytes, inspect::SurfaceBudget::default());
309    if surface.truncated() {
310        tracing::warn!(
311            ai_request_id = %ai_request_id,
312            leaves = surface.len(),
313            "Gateway inspection surface truncated — part of the forwarded body was not scanned"
314        );
315    }
316    request.forwarded_surface = surface;
317}
318
319fn outbound_ctx<'a>(
320    upstream: &'a ResolvedUpstream<'a>,
321    request: &'a CanonicalRequest,
322    parts: CtxParts<'a>,
323) -> OutboundCtx<'a> {
324    OutboundCtx {
325        route: upstream.route.as_ref(),
326        endpoint: &upstream.provider.endpoint,
327        api_key: upstream.api_key,
328        request,
329        upstream_model: parts.upstream_model,
330        model_limits: parts.model_limits,
331        forward_headers: parts.forward_headers,
332        raw_body: parts.raw_body,
333    }
334}
335
336async fn send_to_upstream(
337    upstream: &ResolvedUpstream<'_>,
338    request: &CanonicalRequest,
339    prepared: &Prepared,
340    forward_headers: &[(String, String)],
341    audit: &GatewayAudit,
342) -> Result<OutboundOutcome, DispatchError> {
343    let ctx = outbound_ctx(
344        upstream,
345        request,
346        CtxParts {
347            upstream_model: &prepared.upstream_model,
348            model_limits: prepared.model_limits,
349            forward_headers,
350            raw_body: None,
351        },
352    );
353    match upstream.adapter.send(ctx, &prepared.body).await {
354        Ok(o) => Ok(o),
355        Err(e) => {
356            audit_upstream_failure(audit, upstream.provider.name.as_str(), &request.model, &e)
357                .await;
358            Err(DispatchError::Recorded(e))
359        },
360    }
361}
362
363// Why: `metadata.user_id` is an end-user identifier meant for the provider the
364// caller chose, so it must not reach a different wire's upstream. The
365// passthrough lane applies the same rule to the raw body in
366// `normalize_raw_body`.
367fn strip_caller_identity(request: &mut CanonicalRequest) {
368    let Some(metadata) = request.metadata.as_mut() else {
369        return;
370    };
371    let Some(obj) = metadata.as_object_mut() else {
372        return;
373    };
374    obj.remove("user_id");
375    if obj.is_empty() {
376        request.metadata = None;
377    }
378}
379
380async fn enforce_quota(
381    db: &DbPool,
382    user_id: &UserId,
383    quota_windows: &[QuotaWindow],
384    audit: &GatewayAudit,
385) -> Result<(), DispatchError> {
386    let reservation = quota::precheck_and_reserve(db, 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    db: &DbPool,
600    ai_request_id: &AiRequestId,
601    request: &CanonicalRequest,
602    safety: &SafetyConfig,
603    audit: &GatewayAudit,
604) -> Result<(), DispatchError> {
605    let findings = run_request_safety_scan(db, ai_request_id, request, safety).await;
606    let Some(finding) = findings.iter().find(|f| {
607        safety.block_categories.contains(&f.category) && blocks_at_phase(f.phase, safety.history)
608    }) else {
609        return Ok(());
610    };
611    let msg = format!(
612        "request blocked by safety policy: category '{}'",
613        finding.category
614    );
615    tracing::warn!(
616        ai_request_id = %ai_request_id,
617        category = %finding.category,
618        scanner = %finding.scanner,
619        "Gateway blocked request by safety policy"
620    );
621    if let Err(e) = audit.fail(&msg).await {
622        tracing::warn!(error = %e, "safety-block audit fail failed");
623    }
624    Err(DispatchError::Recorded(
625        SafetyBlocked {
626            category: finding.category.clone(),
627            message: msg,
628        }
629        .into(),
630    ))
631}
632
633/// Whether a finding raised at `phase` may deny the request.
634///
635/// A blocked category found in an earlier turn would otherwise deny every
636/// remaining turn of the conversation, including the turns that carry nothing
637/// objectionable — and a tool call the policy layer already refused is replayed
638/// into the scan surface for the rest of the session.
639pub fn blocks_at_phase(phase: &str, history: SafetyHistoryMode) -> bool {
640    match phase {
641        PHASE_REQUEST => true,
642        PHASE_REQUEST_HISTORY => history == SafetyHistoryMode::Block,
643        _ => false,
644    }
645}
646
647async fn audit_upstream_failure(
648    audit: &GatewayAudit,
649    provider: &str,
650    model: &str,
651    error: &anyhow::Error,
652) {
653    tracing::warn!(
654        provider = %provider,
655        model = %model,
656        error = %error,
657        "gateway upstream call failed"
658    );
659    if let Err(audit_err) = audit.fail(&error.to_string()).await {
660        tracing::warn!(error = %audit_err, "upstream audit fail failed");
661    }
662}
systemprompt_api/services/gateway/service/mod.rs

systemprompt_api/services/gateway/service/
mod.rs
