mockforge_intelligence/ai_studio/
org_controls.rs

1//! Organization-level AI controls service
2//!
3//! This module provides functionality to manage org-level AI controls including
4//! budgets, rate limits, and feature toggles. Supports YAML defaults with DB
5//! authoritative overrides (DB overrides YAML).
6
7use crate::ai_studio::budget_manager::AiFeature;
8use async_trait::async_trait;
9use chrono::{DateTime, Utc};
10use mockforge_foundation::Result;
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13
14/// Organization AI controls configuration
15#[derive(Debug, Clone, Serialize, Deserialize)]
16#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
17pub struct OrgAiControlsConfig {
18    /// Budget configuration
19    pub budgets: OrgBudgetConfig,
20    /// Rate limit configuration
21    pub rate_limits: OrgRateLimitConfig,
22    /// Feature toggles
23    pub feature_toggles: HashMap<String, bool>,
24}
25
26impl Default for OrgAiControlsConfig {
27    fn default() -> Self {
28        Self {
29            budgets: OrgBudgetConfig::default(),
30            rate_limits: OrgRateLimitConfig::default(),
31            feature_toggles: HashMap::from([
32                ("mock_generation".to_string(), true),
33                ("contract_diff".to_string(), true),
34                ("persona_generation".to_string(), true),
35                ("free_form_generation".to_string(), true),
36                ("debug_analysis".to_string(), true),
37            ]),
38        }
39    }
40}
41
42/// Organization budget configuration
43#[derive(Debug, Clone, Serialize, Deserialize)]
44#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
45pub struct OrgBudgetConfig {
46    /// Maximum tokens per period
47    pub max_tokens_per_period: u64,
48    /// Period type (day, week, month, year)
49    pub period_type: String,
50    /// Maximum AI calls per period
51    pub max_calls_per_period: u64,
52}
53
54impl Default for OrgBudgetConfig {
55    fn default() -> Self {
56        Self {
57            max_tokens_per_period: 1_000_000,
58            period_type: "month".to_string(),
59            max_calls_per_period: 10_000,
60        }
61    }
62}
63
64/// Organization rate limit configuration
65#[derive(Debug, Clone, Serialize, Deserialize)]
66#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
67pub struct OrgRateLimitConfig {
68    /// Rate limit per minute
69    pub rate_limit_per_minute: u64,
70    /// Optional rate limit per hour
71    pub rate_limit_per_hour: Option<u64>,
72    /// Optional rate limit per day
73    pub rate_limit_per_day: Option<u64>,
74}
75
76impl Default for OrgRateLimitConfig {
77    fn default() -> Self {
78        Self {
79            rate_limit_per_minute: 100,
80            rate_limit_per_hour: None,
81            rate_limit_per_day: None,
82        }
83    }
84}
85
86/// Trait for accessing org controls from database
87#[allow(clippy::too_many_arguments)]
88#[async_trait]
89pub trait OrgControlsAccessor: Send + Sync {
90    /// Load org controls configuration
91    async fn load_org_config(
92        &self,
93        org_id: &str,
94        workspace_id: Option<&str>,
95    ) -> Result<Option<OrgAiControlsConfig>>;
96
97    /// Check if budget allows the request
98    async fn check_budget(
99        &self,
100        org_id: &str,
101        workspace_id: Option<&str>,
102        estimated_tokens: u64,
103    ) -> Result<BudgetCheckResult>;
104
105    /// Check if rate limit allows the request
106    async fn check_rate_limit(
107        &self,
108        org_id: &str,
109        workspace_id: Option<&str>,
110    ) -> Result<RateLimitCheckResult>;
111
112    /// Check if a feature is enabled
113    async fn is_feature_enabled(
114        &self,
115        org_id: &str,
116        workspace_id: Option<&str>,
117        feature: &str,
118    ) -> Result<bool>;
119
120    /// Record usage for audit
121    async fn record_usage(
122        &self,
123        org_id: &str,
124        workspace_id: Option<&str>,
125        user_id: Option<&str>,
126        feature: AiFeature,
127        tokens: u64,
128        cost_usd: f64,
129        metadata: Option<serde_json::Value>,
130    ) -> Result<()>;
131}
132
133/// Result of budget check
134#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct BudgetCheckResult {
136    /// Whether the request is allowed
137    pub allowed: bool,
138    /// Current tokens used in period
139    pub current_tokens: u64,
140    /// Maximum tokens allowed in period
141    pub max_tokens: u64,
142    /// Current calls used in period
143    pub current_calls: u64,
144    /// Maximum calls allowed in period
145    pub max_calls: u64,
146    /// Period start timestamp
147    pub period_start: Option<DateTime<Utc>>,
148    /// Reason if not allowed
149    pub reason: Option<String>,
150}
151
152/// Result of rate limit check
153#[derive(Debug, Clone, Serialize, Deserialize)]
154pub struct RateLimitCheckResult {
155    /// Whether the request is allowed
156    pub allowed: bool,
157    /// Current requests in the current window
158    pub current_requests: u64,
159    /// Maximum requests allowed in the window
160    pub max_requests: u64,
161    /// Window type (minute, hour, day)
162    pub window_type: String,
163    /// Retry after timestamp (if rate limited)
164    pub retry_after: Option<DateTime<Utc>>,
165    /// Reason if not allowed
166    pub reason: Option<String>,
167}
168
169/// Organization controls service
170///
171/// Manages org-level AI controls with YAML defaults and DB authoritative overrides.
172/// DB values override YAML defaults when both are present.
173pub struct OrgControls {
174    /// YAML-based default configuration
175    yaml_config: OrgAiControlsConfig,
176    /// Optional database accessor (if available)
177    db_accessor: Option<Box<dyn OrgControlsAccessor>>,
178}
179
180impl OrgControls {
181    /// Create a new org controls service with YAML defaults only
182    pub fn new(yaml_config: OrgAiControlsConfig) -> Self {
183        Self {
184            yaml_config,
185            db_accessor: None,
186        }
187    }
188
189    /// Create a new org controls service with YAML defaults and DB accessor
190    pub fn with_db_accessor(
191        yaml_config: OrgAiControlsConfig,
192        db_accessor: Box<dyn OrgControlsAccessor>,
193    ) -> Self {
194        Self {
195            yaml_config,
196            db_accessor: Some(db_accessor),
197        }
198    }
199
200    /// Load org configuration (DB overrides YAML)
201    pub async fn load_org_config(
202        &self,
203        org_id: &str,
204        workspace_id: Option<&str>,
205    ) -> Result<OrgAiControlsConfig> {
206        // Try to load from DB first
207        if let Some(ref accessor) = self.db_accessor {
208            if let Some(db_config) = accessor.load_org_config(org_id, workspace_id).await? {
209                // Merge DB config with YAML defaults (DB values take precedence)
210                return Ok(self.merge_configs(self.yaml_config.clone(), db_config));
211            }
212        }
213
214        // Fall back to YAML config
215        Ok(self.yaml_config.clone())
216    }
217
218    /// Check if budget allows the request
219    pub async fn check_budget(
220        &self,
221        org_id: &str,
222        workspace_id: Option<&str>,
223        estimated_tokens: u64,
224    ) -> Result<BudgetCheckResult> {
225        // Check DB first if available
226        if let Some(ref accessor) = self.db_accessor {
227            return accessor.check_budget(org_id, workspace_id, estimated_tokens).await;
228        }
229
230        // Fall back to YAML config check (simplified - no period tracking)
231        let config = self.load_org_config(org_id, workspace_id).await?;
232        Ok(BudgetCheckResult {
233            allowed: true, // YAML-only mode: always allow (no enforcement)
234            current_tokens: 0,
235            max_tokens: config.budgets.max_tokens_per_period,
236            current_calls: 0,
237            max_calls: config.budgets.max_calls_per_period,
238            period_start: None,
239            reason: None,
240        })
241    }
242
243    /// Check if rate limit allows the request
244    pub async fn check_rate_limit(
245        &self,
246        org_id: &str,
247        workspace_id: Option<&str>,
248    ) -> Result<RateLimitCheckResult> {
249        // Check DB first if available
250        if let Some(ref accessor) = self.db_accessor {
251            return accessor.check_rate_limit(org_id, workspace_id).await;
252        }
253
254        // Fall back to YAML config check (simplified - no rate limiting)
255        let config = self.load_org_config(org_id, workspace_id).await?;
256        Ok(RateLimitCheckResult {
257            allowed: true, // YAML-only mode: always allow (no enforcement)
258            current_requests: 0,
259            max_requests: config.rate_limits.rate_limit_per_minute,
260            window_type: "minute".to_string(),
261            retry_after: None,
262            reason: None,
263        })
264    }
265
266    /// Check if a feature is enabled
267    pub async fn is_feature_enabled(
268        &self,
269        org_id: &str,
270        workspace_id: Option<&str>,
271        feature: &str,
272    ) -> Result<bool> {
273        // Check DB first if available
274        if let Some(ref accessor) = self.db_accessor {
275            return accessor.is_feature_enabled(org_id, workspace_id, feature).await;
276        }
277
278        // Fall back to YAML config
279        let config = self.load_org_config(org_id, workspace_id).await?;
280        Ok(config.feature_toggles.get(feature).copied().unwrap_or(true))
281    }
282
283    /// Record usage for audit
284    #[allow(clippy::too_many_arguments)]
285    pub async fn record_usage(
286        &self,
287        org_id: &str,
288        workspace_id: Option<&str>,
289        user_id: Option<&str>,
290        feature: AiFeature,
291        tokens: u64,
292        cost_usd: f64,
293        metadata: Option<serde_json::Value>,
294    ) -> Result<()> {
295        // Record in DB if available
296        if let Some(ref accessor) = self.db_accessor {
297            return accessor
298                .record_usage(org_id, workspace_id, user_id, feature, tokens, cost_usd, metadata)
299                .await;
300        }
301
302        // YAML-only mode: no recording (in-memory tracking would be handled by BudgetManager)
303        Ok(())
304    }
305
306    /// Merge DB config with YAML defaults (DB values take precedence)
307    fn merge_configs(
308        &self,
309        yaml: OrgAiControlsConfig,
310        db: OrgAiControlsConfig,
311    ) -> OrgAiControlsConfig {
312        // Merge feature toggles (DB overrides YAML)
313        let mut merged_toggles = yaml.feature_toggles.clone();
314        for (key, value) in db.feature_toggles {
315            merged_toggles.insert(key, value);
316        }
317
318        OrgAiControlsConfig {
319            budgets: db.budgets,         // DB budget config takes precedence
320            rate_limits: db.rate_limits, // DB rate limit config takes precedence
321            feature_toggles: merged_toggles,
322        }
323    }
324}
325
326impl Default for OrgControls {
327    fn default() -> Self {
328        Self::new(OrgAiControlsConfig::default())
329    }
330}
mockforge_intelligence/ai_studio/org_controls.rs

mockforge_intelligence/ai_studio/
org_controls.rs
