Skip to main content

agentic_contract/
risk_limit.rs

1//! Risk limit thresholds for agent actions.
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5
6use crate::ContractId;
7
8/// Type of risk limit.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum LimitType {
12    /// Rate limit (actions per time window).
13    Rate,
14    /// Threshold limit (value must stay below).
15    Threshold,
16    /// Budget limit (cumulative spending cap).
17    Budget,
18    /// Count limit (total number of actions).
19    Count,
20}
21
22/// A risk limit threshold for a resource or action.
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct RiskLimit {
25    /// Unique identifier.
26    pub id: ContractId,
27    /// Human-readable label.
28    pub label: String,
29    /// Type of limit.
30    pub limit_type: LimitType,
31    /// Current accumulated value.
32    pub current_value: f64,
33    /// Maximum allowed value.
34    pub max_value: f64,
35    /// Time window in seconds (for rate limits).
36    #[serde(default)]
37    pub window_secs: Option<u64>,
38    /// When the current window started.
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub window_start: Option<DateTime<Utc>>,
41    /// When this limit was created.
42    pub created_at: DateTime<Utc>,
43    /// When this limit was last updated.
44    pub updated_at: DateTime<Utc>,
45}
46
47impl RiskLimit {
48    /// Create a new risk limit.
49    pub fn new(label: impl Into<String>, limit_type: LimitType, max_value: f64) -> Self {
50        let now = Utc::now();
51        Self {
52            id: ContractId::new(),
53            label: label.into(),
54            limit_type,
55            current_value: 0.0,
56            max_value,
57            window_secs: None,
58            window_start: None,
59            created_at: now,
60            updated_at: now,
61        }
62    }
63
64    /// Set a time window (for rate limits).
65    pub fn with_window(mut self, secs: u64) -> Self {
66        self.window_secs = Some(secs);
67        self.window_start = Some(Utc::now());
68        self
69    }
70
71    /// Check if the limit would be exceeded by adding `amount`.
72    pub fn would_exceed(&self, amount: f64) -> bool {
73        self.current_value + amount > self.max_value
74    }
75
76    /// Get remaining capacity.
77    pub fn remaining(&self) -> f64 {
78        (self.max_value - self.current_value).max(0.0)
79    }
80
81    /// Get usage as a percentage (0.0 to 1.0).
82    pub fn usage_ratio(&self) -> f64 {
83        if self.max_value == 0.0 {
84            return 1.0;
85        }
86        (self.current_value / self.max_value).min(1.0)
87    }
88
89    /// Increment the current value.
90    pub fn increment(&mut self, amount: f64) {
91        self.current_value += amount;
92        self.updated_at = Utc::now();
93    }
94
95    /// Reset the current value (e.g. for a new time window).
96    pub fn reset(&mut self) {
97        self.current_value = 0.0;
98        self.window_start = Some(Utc::now());
99        self.updated_at = Utc::now();
100    }
101
102    /// Check if the time window has expired and needs resetting.
103    pub fn window_expired(&self) -> bool {
104        if let (Some(window_secs), Some(window_start)) = (self.window_secs, self.window_start) {
105            let elapsed = (Utc::now() - window_start).num_seconds() as u64;
106            elapsed >= window_secs
107        } else {
108            false
109        }
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn test_risk_limit_creation() {
119        let limit = RiskLimit::new("API calls per minute", LimitType::Rate, 100.0).with_window(60);
120
121        assert_eq!(limit.label, "API calls per minute");
122        assert_eq!(limit.limit_type, LimitType::Rate);
123        assert_eq!(limit.max_value, 100.0);
124        assert_eq!(limit.current_value, 0.0);
125        assert!(limit.window_secs.is_some());
126    }
127
128    #[test]
129    fn test_would_exceed() {
130        let mut limit = RiskLimit::new("Budget", LimitType::Budget, 1000.0);
131        limit.increment(900.0);
132
133        assert!(!limit.would_exceed(50.0));
134        assert!(limit.would_exceed(200.0));
135        assert_eq!(limit.remaining(), 100.0);
136    }
137
138    #[test]
139    fn test_usage_ratio() {
140        let mut limit = RiskLimit::new("Count", LimitType::Count, 10.0);
141        limit.increment(5.0);
142        assert!((limit.usage_ratio() - 0.5).abs() < f64::EPSILON);
143    }
144}