ares_agent/
admit.rs

1use std::sync::Arc;
2
3use ares_types::models::{QuotaExceeded, TenantContext};
4use ares_types::types::AppError;
5use cordis::{Context, CordisError, EventsService};
6use serde_json::Value;
7
8/// Which usage query failed while preparing the admission payload.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum UsagePeriod {
11    Monthly,
12    Daily,
13}
14
15/// Failure from the shared admission gate.
16#[derive(Debug)]
17pub enum AdmissionError {
18    Usage {
19        period: UsagePeriod,
20        source: AppError,
21    },
22    Event(CordisError),
23    Quota(QuotaExceeded),
24}
25
26impl From<AdmissionError> for AppError {
27    fn from(error: AdmissionError) -> Self {
28        match error {
29            AdmissionError::Usage { source, .. } => source,
30            AdmissionError::Event(error) => {
31                AppError::Internal(format!("admission event failed: {error}"))
32            }
33            AdmissionError::Quota(exceeded) => exceeded.into(),
34        }
35    }
36}
37
38/// Apply the final typed quota policy to a usage snapshot.
39pub fn quota_exceeded(tenant: &TenantContext, monthly: u64, daily: u64) -> Option<QuotaExceeded> {
40    tenant.admit(monthly, daily).err()
41}
42
43/// Shared quota gate used by `Execute::run` and protocol adapters.
44///
45/// The event is authoritative when an `EventsService` is available. The typed
46/// `TenantContext::admit` check remains the final fallback, which keeps direct
47/// library contexts safe when no event bus has been installed yet.
48pub async fn admit(ctx: &Arc<Context>) -> Result<(), AppError> {
49    admit_with_details(ctx).await.map_err(Into::into)
50}
51
52/// Shared admission gate with enough detail for protocol-specific error maps.
53pub async fn admit_with_details(ctx: &Arc<Context>) -> Result<(), AdmissionError> {
54    let Some(tc) = ctx.get::<TenantContext>() else {
55        return Ok(());
56    };
57    let (monthly, daily) = usage_counts(ctx, &tc.tenant_id).await?;
58    if let Some(events) = ctx.get::<EventsService>() {
59        let payload = cordis::AgentAdmitPayload {
60            tenant_id: tc.tenant_id.clone(),
61            monthly,
62            daily,
63            requests_per_month: Some(tc.quota.requests_per_month),
64            requests_per_day: Some(tc.quota.requests_per_day),
65            tier: tc.tier.as_str().to_string(),
66        };
67        let result = events
68            .dispatch_typed::<cordis::AgentAdmitEvent>(&payload)
69            .await
70            .map_err(AdmissionError::Event)?;
71        if let Some(err) = deny_from_bail(&result) {
72            return Err(AdmissionError::Quota(err));
73        }
74    }
75    quota_exceeded(&tc, monthly, daily)
76        .map_or(Ok(()), |exceeded| Err(AdmissionError::Quota(exceeded)))
77}
78
79fn deny_from_bail(result: &Value) -> Option<QuotaExceeded> {
80    let marker = result
81        .get("deny")
82        .and_then(|v| v.as_str())
83        .or_else(|| result.get("error").and_then(|v| v.as_str()));
84    match marker {
85        Some("daily") => Some(QuotaExceeded::Daily),
86        Some("monthly") | Some(_) => Some(QuotaExceeded::Monthly),
87        None => None,
88    }
89}
90
91async fn usage_counts(ctx: &Arc<Context>, tenant_id: &str) -> Result<(u64, u64), AdmissionError> {
92    #[cfg(feature = "postgres")]
93    {
94        if let Some(db) = ctx.get::<ares_store::TenantDb>() {
95            let monthly = db.get_monthly_requests(tenant_id).await.map_err(|source| {
96                AdmissionError::Usage {
97                    period: UsagePeriod::Monthly,
98                    source,
99                }
100            })?;
101            let daily =
102                db.get_daily_requests(tenant_id)
103                    .await
104                    .map_err(|source| AdmissionError::Usage {
105                        period: UsagePeriod::Daily,
106                        source,
107                    })?;
108            return Ok((monthly, daily));
109        }
110    }
111    let _ = (ctx, tenant_id);
112    Ok((0, 0))
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use ares_types::models::TenantTier;
119    use serde_json::json;
120
121    fn free_tenant() -> TenantContext {
122        TenantContext::new("acme".into(), TenantTier::Free)
123    }
124
125    fn ctx_with_deny(deny: &'static str) -> (Arc<Context>, Box<dyn cordis::Disposable>) {
126        let root = Context::new_root();
127        let events = root.provide(EventsService::new());
128        let keep = events.on(
129            cordis::events_catalog::ev::AGENT_ADMIT.to_string(),
130            move |_payload| async move { Ok(json!({ "deny": deny })) },
131        );
132        let ctx = root.with_intercept(free_tenant());
133        (ctx, keep)
134    }
135
136    #[tokio::test]
137    async fn bail_deny_monthly_overrides_passing_typed_quota() {
138        let tenant = free_tenant();
139        assert!(
140            tenant.admit(0, 0).is_ok(),
141            "typed Free quota must pass at zero usage"
142        );
143        let (ctx, _keep) = ctx_with_deny("monthly");
144        let err = admit_with_details(&ctx)
145            .await
146            .expect_err("event deny must win over typed pass");
147        assert!(matches!(err, AdmissionError::Quota(QuotaExceeded::Monthly)));
148    }
149
150    #[tokio::test]
151    async fn bail_deny_daily_overrides_passing_typed_quota() {
152        let tenant = free_tenant();
153        assert!(
154            tenant.admit(0, 0).is_ok(),
155            "typed Free quota must pass at zero usage"
156        );
157        let (ctx, _keep) = ctx_with_deny("daily");
158        let err = admit_with_details(&ctx)
159            .await
160            .expect_err("event deny must win over typed pass");
161        assert!(matches!(err, AdmissionError::Quota(QuotaExceeded::Daily)));
162    }
163
164    #[tokio::test]
165    async fn admit_without_events_uses_typed_fallback() {
166        let ctx = Context::new_root().with_intercept(free_tenant());
167        assert!(
168            ctx.get::<EventsService>().is_none(),
169            "this path must not install EventsService"
170        );
171        assert!(quota_exceeded(&free_tenant(), 0, 0).is_none());
172        admit_with_details(&ctx)
173            .await
174            .expect("typed fallback admits Free quota at zero usage");
175    }
176}
ares_agent/admit.rs

ares_agent/
admit.rs
