Skip to main content

xcom_rs/
context.rs

1use crate::billing::{BudgetTracker, CostEstimate};
2use crate::protocol::{ErrorCode, ErrorDetails};
3
4/// Execution context for commands
5#[derive(Debug, Clone)]
6pub struct ExecutionContext {
7    /// Whether running in non-interactive mode
8    pub non_interactive: bool,
9    /// Optional trace ID for request correlation
10    pub trace_id: Option<String>,
11    /// Maximum cost in credits for a single operation
12    pub max_cost_credits: Option<u32>,
13    /// Daily budget in credits
14    pub budget_daily_credits: Option<u32>,
15    /// Dry run mode
16    pub dry_run: bool,
17}
18
19impl ExecutionContext {
20    /// Create a new execution context
21    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    /// Check if authentication is required and return an error if in non-interactive mode
38    ///
39    /// This helper should be called by commands that need authentication or user interaction.
40    /// If in non-interactive mode, it returns an AUTH_REQUIRED error with next steps.
41    /// Otherwise, it allows the command to proceed with interactive prompts.
42    ///
43    /// # Example
44    /// ```
45    /// use xcom_rs::context::ExecutionContext;
46    ///
47    /// let ctx = ExecutionContext::new(true, None, None, None, false);
48    /// let error = ctx.check_interaction_required(
49    ///     "Authentication required",
50    ///     vec!["Run 'xcom-rs auth login' first".to_string()]
51    /// );
52    /// // If error.is_some(), handle the interaction requirement
53    /// assert!(error.is_some());
54    /// ```
55    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    /// Check if cost exceeds maximum allowed
68    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    /// Check if cost would exceed daily budget
88    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}