Struct BudgetManager 

pub struct BudgetManager { /* private fields */ }

Budget manager for AI usage

Implementations

impl BudgetManager

pub fn new(config: BudgetConfig) -> Self

Create a new budget manager

pub fn with_org_controls( config: BudgetConfig, org_controls: Arc<OrgControls>, ) -> Self

Create a new budget manager with org controls

pub async fn get_usage(&self, workspace_id: &str) -> Result<UsageStats>

Get usage statistics for a workspace

pub async fn check_budget( &self, org_id: Option<&str>, workspace_id: &str, estimated_tokens: u64, ) -> Result<bool>

Check if request is within budget

Checks org-level limits first (if available), then workspace-level limits. Org-level limits take precedence.

pub async fn check_rate_limit( &self, org_id: Option<&str>, workspace_id: &str, ) -> Result<bool>

Check rate limit (org-level first, then workspace-level)

pub async fn is_feature_enabled( &self, org_id: Option<&str>, workspace_id: &str, feature: &str, ) -> Result<bool>

Check if a feature is enabled (org-level first, then defaults to true)

pub async fn record_usage( &self, org_id: Option<&str>, workspace_id: &str, user_id: Option<&str>, tokens: u64, cost_usd: f64, ) -> Result<()>

Record token usage and cost

pub async fn record_usage_with_feature( &self, org_id: Option<&str>, workspace_id: &str, user_id: Option<&str>, tokens: u64, cost_usd: f64, feature: Option<AiFeature>, ) -> Result<()>

Record token usage and cost with feature tracking

Records usage both in-memory (workspace-level) and in org controls (if available).

pub async fn reset_usage(&self, workspace_id: &str) -> Result<()>

Reset usage for a workspace (useful for testing or monthly resets)

pub fn calculate_cost(provider: &str, model: &str, tokens: u64) -> f64

Calculate cost based on provider and tokens

Uses approximate pricing for common providers:

• OpenAI GPT-3.5: ~$0.002 per 1K tokens
• OpenAI GPT-4: ~$0.03 per 1K tokens
• Anthropic Claude: ~$0.008 per 1K tokens
• Ollama: $0 (local)