agentic_contract/
risk_limit.rs1use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5
6use crate::ContractId;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum LimitType {
12 Rate,
14 Threshold,
16 Budget,
18 Count,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct RiskLimit {
25 pub id: ContractId,
27 pub label: String,
29 pub limit_type: LimitType,
31 pub current_value: f64,
33 pub max_value: f64,
35 #[serde(default)]
37 pub window_secs: Option<u64>,
38 #[serde(skip_serializing_if = "Option::is_none")]
40 pub window_start: Option<DateTime<Utc>>,
41 pub created_at: DateTime<Utc>,
43 pub updated_at: DateTime<Utc>,
45}
46
47impl RiskLimit {
48 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 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 pub fn would_exceed(&self, amount: f64) -> bool {
73 self.current_value + amount > self.max_value
74 }
75
76 pub fn remaining(&self) -> f64 {
78 (self.max_value - self.current_value).max(0.0)
79 }
80
81 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 pub fn increment(&mut self, amount: f64) {
91 self.current_value += amount;
92 self.updated_at = Utc::now();
93 }
94
95 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 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}