systemprompt_api/services/gateway/
policy.rs

1use std::collections::HashMap;
2use std::sync::{Arc, RwLock};
3use std::time::{Duration, Instant};
4
5use anyhow::Result;
6use serde::{Deserialize, Serialize};
7use systemprompt_ai::repository::AiGatewayPolicyRepository;
8use systemprompt_database::DbPool;
9use systemprompt_identifiers::TenantId;
10
11#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
12pub struct QuotaWindow {
13    pub window_seconds: i32,
14    pub max_requests: Option<i64>,
15    pub max_input_tokens: Option<i64>,
16    pub max_output_tokens: Option<i64>,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize, Default)]
20pub struct SafetyConfig {
21    #[serde(default)]
22    pub scanners: Vec<String>,
23    #[serde(default)]
24    pub block_categories: Vec<String>,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize, Default)]
28pub struct GatewayPolicySpec {
29    #[serde(default)]
30    pub allowed_models: Option<Vec<String>>,
31    #[serde(default)]
32    pub max_input_tokens_per_call: Option<u32>,
33    #[serde(default)]
34    pub max_tool_depth: Option<u32>,
35    #[serde(default)]
36    pub quota_windows: Vec<QuotaWindow>,
37    #[serde(default)]
38    pub safety: SafetyConfig,
39}
40
41impl GatewayPolicySpec {
42    pub fn permissive() -> Self {
43        Self::default()
44    }
45
46    pub fn model_allowed(&self, model: &str) -> bool {
47        self.allowed_models
48            .as_deref()
49            .is_none_or(|list| list.iter().any(|m| m == model))
50    }
51}
52
53const CACHE_TTL: Duration = Duration::from_secs(60);
54
55#[derive(Clone)]
56pub struct PolicyResolver {
57    repo: Arc<AiGatewayPolicyRepository>,
58    cache: Arc<RwLock<HashMap<String, CachedEntry>>>,
59}
60
61impl std::fmt::Debug for PolicyResolver {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        f.debug_struct("PolicyResolver").finish()
64    }
65}
66
67#[derive(Clone)]
68struct CachedEntry {
69    spec: GatewayPolicySpec,
70    fetched_at: Instant,
71}
72
73impl PolicyResolver {
74    pub fn new(db: &DbPool) -> Result<Self> {
75        Ok(Self {
76            repo: Arc::new(
77                AiGatewayPolicyRepository::new(db)
78                    .map_err(|e| anyhow::anyhow!("policy repo init: {e}"))?,
79            ),
80            cache: Arc::new(RwLock::new(HashMap::new())),
81        })
82    }
83
84    pub async fn resolve(&self, tenant_id: Option<&TenantId>) -> GatewayPolicySpec {
85        let key = tenant_id.map_or_else(String::new, |t| t.as_str().to_string());
86
87        if let Ok(cache) = self.cache.read() {
88            if let Some(entry) = cache.get(&key) {
89                if entry.fetched_at.elapsed() < CACHE_TTL {
90                    return entry.spec.clone();
91                }
92            }
93        }
94
95        let rows = match self.repo.find_for_tenant(tenant_id).await {
96            Ok(r) => r,
97            Err(e) => {
98                tracing::warn!(error = %e, "policy resolve DB error — falling back to permissive");
99                return GatewayPolicySpec::permissive();
100            },
101        };
102
103        let spec = merge(rows);
104        if let Ok(mut cache) = self.cache.write() {
105            cache.insert(
106                key,
107                CachedEntry {
108                    spec: spec.clone(),
109                    fetched_at: Instant::now(),
110                },
111            );
112        }
113        spec
114    }
115}
116
117fn merge(rows: Vec<systemprompt_ai::GatewayPolicyRow>) -> GatewayPolicySpec {
118    let mut merged = GatewayPolicySpec::permissive();
119    for row in rows {
120        let Ok(spec) = serde_json::from_value::<GatewayPolicySpec>(row.spec) else {
121            tracing::warn!(policy_id = %row.id, name = %row.name, "policy spec JSON malformed — skipped");
122            continue;
123        };
124        if spec.allowed_models.is_some() {
125            merged.allowed_models = spec.allowed_models;
126        }
127        if spec.max_input_tokens_per_call.is_some() {
128            merged.max_input_tokens_per_call = spec.max_input_tokens_per_call;
129        }
130        if spec.max_tool_depth.is_some() {
131            merged.max_tool_depth = spec.max_tool_depth;
132        }
133        if !spec.quota_windows.is_empty() {
134            merged.quota_windows = spec.quota_windows;
135        }
136        if !spec.safety.scanners.is_empty() || !spec.safety.block_categories.is_empty() {
137            merged.safety = spec.safety;
138        }
139    }
140    merged
141}
systemprompt_api/services/gateway/policy.rs

systemprompt_api/services/gateway/
policy.rs
