1use crate::billing::{BudgetTracker, CostEstimate};
2use crate::protocol::{ErrorCode, ErrorDetails};
3
4#[derive(Debug, Clone)]
6pub struct ExecutionContext {
7 pub non_interactive: bool,
9 pub trace_id: Option<String>,
11 pub max_cost_credits: Option<u32>,
13 pub budget_daily_credits: Option<u32>,
15 pub dry_run: bool,
17}
18
19impl ExecutionContext {
20 pub fn new(
22 non_interactive: bool,
23 trace_id: Option<String>,
24 max_cost_credits: Option<u32>,
25 budget_daily_credits: Option<u32>,
26 dry_run: bool,
27 ) -> Self {
28 Self {
29 non_interactive,
30 trace_id,
31 max_cost_credits,
32 budget_daily_credits,
33 dry_run,
34 }
35 }
36
37 pub fn check_interaction_required(
56 &self,
57 message: impl Into<String>,
58 next_steps: Vec<String>,
59 ) -> Option<ErrorDetails> {
60 if self.non_interactive {
61 Some(ErrorDetails::auth_required(message, next_steps))
62 } else {
63 None
64 }
65 }
66
67 pub fn check_max_cost(&self, cost: &CostEstimate) -> Option<ErrorDetails> {
69 if let Some(max) = self.max_cost_credits {
70 if cost.credits > max {
71 let mut details = std::collections::HashMap::new();
72 details.insert("cost".to_string(), serde_json::json!(cost.credits));
73 details.insert("limit".to_string(), serde_json::json!(max));
74 return Some(ErrorDetails::with_details(
75 ErrorCode::CostLimitExceeded,
76 format!(
77 "Operation cost {} credits exceeds maximum {} credits",
78 cost.credits, max
79 ),
80 details,
81 ));
82 }
83 }
84 None
85 }
86
87 pub fn check_daily_budget(
89 &self,
90 cost: &CostEstimate,
91 tracker: &BudgetTracker,
92 ) -> Option<ErrorDetails> {
93 if tracker.check_budget(cost.credits).is_err() {
94 let mut details = std::collections::HashMap::new();
95 details.insert("cost".to_string(), serde_json::json!(cost.credits));
96 details.insert(
97 "todayUsage".to_string(),
98 serde_json::json!(tracker.today_usage()),
99 );
100 if let Some(limit) = self.budget_daily_credits {
101 details.insert("dailyLimit".to_string(), serde_json::json!(limit));
102 }
103 return Some(ErrorDetails::with_details(
104 ErrorCode::DailyBudgetExceeded,
105 "Daily budget exceeded".to_string(),
106 details,
107 ));
108 }
109 None
110 }
111}
112
113#[cfg(test)]
114mod tests {
115 use super::*;
116 use crate::protocol::ErrorCode;
117
118 #[test]
119 fn test_context_creation() {
120 let ctx = ExecutionContext::new(true, Some("trace-123".to_string()), None, None, false);
121 assert!(ctx.non_interactive);
122 assert_eq!(ctx.trace_id, Some("trace-123".to_string()));
123 }
124
125 #[test]
126 fn test_check_interaction_required_non_interactive() {
127 let ctx = ExecutionContext::new(true, None, None, None, false);
128 let error =
129 ctx.check_interaction_required("Auth required", vec!["Run login command".to_string()]);
130 assert!(error.is_some());
131 let err = error.unwrap();
132 assert_eq!(err.code, ErrorCode::AuthRequired);
133 assert!(!err.is_retryable);
134 assert!(err.details.is_some());
135 }
136
137 #[test]
138 fn test_check_interaction_required_interactive() {
139 let ctx = ExecutionContext::new(false, None, None, None, false);
140 let error =
141 ctx.check_interaction_required("Auth required", vec!["Run login command".to_string()]);
142 assert!(error.is_none());
143 }
144
145 #[test]
146 fn test_check_max_cost_within_limit() {
147 let ctx = ExecutionContext::new(false, None, Some(100), None, false);
148 let cost = CostEstimate::new(50, 0.05);
149 let error = ctx.check_max_cost(&cost);
150 assert!(error.is_none());
151 }
152
153 #[test]
154 fn test_check_max_cost_exceeds_limit() {
155 let ctx = ExecutionContext::new(false, None, Some(100), None, false);
156 let cost = CostEstimate::new(101, 0.101);
157 let error = ctx.check_max_cost(&cost);
158 assert!(error.is_some());
159 let err = error.unwrap();
160 assert_eq!(err.code, ErrorCode::CostLimitExceeded);
161 }
162
163 #[test]
164 fn test_check_daily_budget_within_limit() {
165 let ctx = ExecutionContext::new(false, None, None, Some(100), false);
166 let tracker = BudgetTracker::new(Some(100));
167 let cost = CostEstimate::new(50, 0.05);
168 let error = ctx.check_daily_budget(&cost, &tracker);
169 assert!(error.is_none());
170 }
171
172 #[test]
173 fn test_check_daily_budget_exceeds_limit() {
174 let ctx = ExecutionContext::new(false, None, None, Some(100), false);
175 let tracker = BudgetTracker::new(Some(100));
176 let cost = CostEstimate::new(101, 0.101);
177 let error = ctx.check_daily_budget(&cost, &tracker);
178 assert!(error.is_some());
179 let err = error.unwrap();
180 assert_eq!(err.code, ErrorCode::DailyBudgetExceeded);
181 }
182}