systemprompt_api/services/gateway/
quota.rs

1//! Quota-window alignment and bucket accounting for gateway policies.
2//!
3//! Windows are keyed by a subject: the requesting user by default, or any
4//! subject-attribute dimension an extension registers (for example
5//! `organization`). Cost ceilings are enforced one request late — cost is
6//! known only after the response.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11use std::sync::{Arc, OnceLock};
12
13use anyhow::Result;
14use chrono::{DateTime, TimeZone, Utc};
15use sqlx::PgPool;
16use systemprompt_ai::USER_QUOTA_SUBJECT;
17use systemprompt_ai::repository::{
18    AiQuotaBucketRepository, IncrementParams, QuotaBucketDelta, QuotaBucketState,
19};
20use systemprompt_database::DbPool;
21use systemprompt_identifiers::UserId;
22use systemprompt_security::authz::{
23    AuthzHookContext, NullAuditSink, SharedSubjectAttributeProvider, discover_subject_providers,
24};
25
26use super::policy::QuotaWindow;
27
28#[derive(Debug, Clone)]
29pub struct QuotaDecision {
30    pub allow: bool,
31    pub window_seconds: i32,
32    pub message: String,
33    pub state: QuotaBucketState,
34}
35
36struct WindowSubject<'a> {
37    kind: &'a str,
38    id: String,
39}
40
41fn subject_providers(pool: &Arc<PgPool>) -> &'static [SharedSubjectAttributeProvider] {
42    static PROVIDERS: OnceLock<Vec<SharedSubjectAttributeProvider>> = OnceLock::new();
43    PROVIDERS.get_or_init(|| {
44        discover_subject_providers(&AuthzHookContext {
45            pool: Arc::clone(pool),
46            sink: Arc::new(NullAuditSink),
47        })
48    })
49}
50
51async fn resolve_subject<'a>(
52    window: &'a QuotaWindow,
53    user_id: &UserId,
54    pool: &Arc<PgPool>,
55) -> Option<WindowSubject<'a>> {
56    if window.subject == USER_QUOTA_SUBJECT {
57        return Some(WindowSubject {
58            kind: USER_QUOTA_SUBJECT,
59            id: user_id.as_str().to_owned(),
60        });
61    }
62    let provider = subject_providers(pool)
63        .iter()
64        .find(|p| p.dimension().rule_type.as_str() == window.subject)?;
65    let id = provider.values_for(user_id).await.into_iter().next()?;
66    Some(WindowSubject {
67        kind: &window.subject,
68        id,
69    })
70}
71
72pub async fn precheck_and_reserve(
73    db: &DbPool,
74    repo: &AiQuotaBucketRepository,
75    user_id: &UserId,
76    windows: &[QuotaWindow],
77) -> Result<Option<QuotaDecision>> {
78    if windows.is_empty() {
79        return Ok(None);
80    }
81    let pool = db
82        .pool_arc()
83        .map_err(|e| anyhow::anyhow!("quota pool init: {e}"))?;
84
85    let now = Utc::now();
86    for window in windows {
87        let Some(subject) = resolve_subject(window, user_id, &pool).await else {
88            continue;
89        };
90        let window_start = align_window(now, window.window_seconds);
91        let state = repo
92            .increment(IncrementParams {
93                subject_kind: subject.kind,
94                subject_id: &subject.id,
95                window_seconds: window.window_seconds,
96                window_start,
97                delta: QuotaBucketDelta {
98                    requests: 1,
99                    input_tokens: 0,
100                    output_tokens: 0,
101                    cost_microdollars: 0,
102                },
103            })
104            .await?;
105
106        if let Some(max) = window.max_requests
107            && state.requests > max
108        {
109            return Ok(Some(QuotaDecision {
110                allow: false,
111                window_seconds: window.window_seconds,
112                message: format!(
113                    "quota exceeded for {} window {}s (used {}/{max})",
114                    subject.kind, window.window_seconds, state.requests
115                ),
116                state,
117            }));
118        }
119
120        if let Some(max) = window.max_cost_microdollars
121            && state.cost_microdollars > max
122        {
123            return Ok(Some(QuotaDecision {
124                allow: false,
125                window_seconds: window.window_seconds,
126                message: format!(
127                    "cost ceiling exceeded for {} window {}s (spent {}/{max} microdollars)",
128                    subject.kind, window.window_seconds, state.cost_microdollars
129                ),
130                state,
131            }));
132        }
133    }
134    Ok(None)
135}
136
137#[derive(Debug)]
138pub struct PostUpdateParams<'a> {
139    pub user_id: &'a UserId,
140    pub windows: &'a [QuotaWindow],
141    pub input_tokens: u32,
142    pub output_tokens: u32,
143    pub cost_microdollars: i64,
144}
145
146pub async fn post_update_tokens(
147    db: &DbPool,
148    repo: &AiQuotaBucketRepository,
149    params: PostUpdateParams<'_>,
150) {
151    if params.windows.is_empty() {
152        return;
153    }
154    let pool = match db.pool_arc() {
155        Ok(p) => p,
156        Err(e) => {
157            tracing::warn!(error = %e, "quota pool init failed in post_update");
158            return;
159        },
160    };
161    let now = Utc::now();
162    for window in params.windows {
163        let Some(subject) = resolve_subject(window, params.user_id, &pool).await else {
164            continue;
165        };
166        let window_start = align_window(now, window.window_seconds);
167        if let Err(e) = repo
168            .increment(IncrementParams {
169                subject_kind: subject.kind,
170                subject_id: &subject.id,
171                window_seconds: window.window_seconds,
172                window_start,
173                delta: QuotaBucketDelta {
174                    requests: 0,
175                    input_tokens: i64::from(params.input_tokens),
176                    output_tokens: i64::from(params.output_tokens),
177                    cost_microdollars: params.cost_microdollars,
178                },
179            })
180            .await
181        {
182            tracing::warn!(error = %e, window_seconds = window.window_seconds, "quota post_update failed");
183        }
184    }
185}
186
187fn align_window(now: DateTime<Utc>, window_seconds: i32) -> DateTime<Utc> {
188    let secs = now.timestamp();
189    let w = i64::from(window_seconds.max(1));
190    let aligned = (secs / w) * w;
191    Utc.timestamp_opt(aligned, 0).single().unwrap_or(now)
192}
systemprompt_api/services/gateway/quota.rs

systemprompt_api/services/gateway/
quota.rs
