1use crate::Cookie;
2use crate::config::AgentConfig;
3use chrono::{DateTime, Utc};
4use serde::Serialize;
5
6pub struct Agent {
7 pub config: AgentConfig,
8 pub cookies: Vec<Cookie>,
9}
10
11#[derive(Debug)]
12pub enum AgentLimit {
13 NotLimited,
14 Limited { reset_time: Option<DateTime<Utc>> },
15}
16
17#[derive(Debug, Serialize)]
18pub struct UsageEntry {
19 #[serde(rename = "type")]
20 pub entry_type: String,
21 pub limited: bool,
22 pub utilization: f64,
23 pub resets_at: Option<DateTime<Utc>>,
24}
25
26#[derive(Debug, Serialize)]
27pub struct AgentStatus {
28 pub command: String,
29 pub provider: Option<String>,
30 pub usage: Vec<UsageEntry>,
31}
32
33fn codex_usage_entries(prefix: &str, limit: &crate::codex::CodexRateLimit) -> Vec<UsageEntry> {
34 let has_limited_window = [
35 limit.primary_window.as_ref(),
36 limit.secondary_window.as_ref(),
37 ]
38 .into_iter()
39 .flatten()
40 .any(|window| window.is_limited());
41 let fallback_reset = if limit.is_limited() && !has_limited_window {
42 limit.next_reset_time()
43 } else {
44 None
45 };
46
47 let mut entries = Vec::new();
48
49 for (suffix, window) in [
50 ("primary", limit.primary_window.as_ref()),
51 ("secondary", limit.secondary_window.as_ref()),
52 ] {
53 if let Some(window) = window {
54 let resets_at = window.reset_at_datetime();
55 entries.push(UsageEntry {
56 entry_type: format!("{}_{}", prefix, suffix),
57 limited: window.is_limited()
58 || (fallback_reset.is_some() && resets_at == fallback_reset),
59 utilization: window.used_percent,
60 resets_at,
61 });
62 }
63 }
64
65 if entries.is_empty() && limit.is_limited() {
66 entries.push(UsageEntry {
67 entry_type: prefix.to_string(),
68 limited: true,
69 utilization: 100.0,
70 resets_at: limit.next_reset_time(),
71 });
72 }
73
74 entries
75}
76
77impl Agent {
78 pub fn new(config: AgentConfig, cookies: Vec<Cookie>) -> Self {
79 Self { config, cookies }
80 }
81
82 pub fn command(&self) -> &str {
83 &self.config.command
84 }
85
86 pub fn args(&self) -> &[String] {
87 &self.config.args
88 }
89
90 pub async fn check_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
91 match self.config.resolve_domain() {
92 Some("claude.ai") => self.check_claude_limit().await,
93 Some("chatgpt.com") => self.check_codex_limit().await,
94 Some("github.com") => self.check_copilot_limit().await,
95 None => Ok(AgentLimit::NotLimited),
96 Some(d) => Err(format!("Unknown domain: {}", d).into()),
97 }
98 }
99
100 pub async fn fetch_status(&self) -> Result<AgentStatus, Box<dyn std::error::Error>> {
101 let command = self.config.command.clone();
102 let provider = self.config.resolve_provider().map(|s| s.to_string());
103 let usage = match self.config.resolve_domain() {
104 None => vec![],
105 Some("claude.ai") => {
106 let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
107 let windows = [
108 ("five_hour", &usage.five_hour),
109 ("seven_day", &usage.seven_day),
110 ("seven_day_sonnet", &usage.seven_day_sonnet),
111 ];
112 windows
113 .into_iter()
114 .filter_map(|(name, w)| {
115 w.as_ref().map(|w| UsageEntry {
116 entry_type: name.to_string(),
117 limited: w.utilization >= 100.0,
118 utilization: w.utilization,
119 resets_at: w.resets_at,
120 })
121 })
122 .collect()
123 }
124 Some("chatgpt.com") => {
125 let usage = crate::codex::CodexClient::fetch_usage(&self.cookies).await?;
126 [
127 ("rate_limit", &usage.rate_limit),
128 ("code_review_rate_limit", &usage.code_review_rate_limit),
129 ]
130 .into_iter()
131 .flat_map(|(prefix, limit)| codex_usage_entries(prefix, limit))
132 .collect()
133 }
134 Some("github.com") => {
135 let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
136 vec![
137 UsageEntry {
138 entry_type: "chat_utilization".to_string(),
139 limited: quota.chat_utilization >= 100.0,
140 utilization: quota.chat_utilization,
141 resets_at: quota.reset_time,
142 },
143 UsageEntry {
144 entry_type: "premium_utilization".to_string(),
145 limited: quota.premium_utilization >= 100.0,
146 utilization: quota.premium_utilization,
147 resets_at: quota.reset_time,
148 },
149 ]
150 }
151 Some(d) => return Err(format!("Unknown domain: {}", d).into()),
152 };
153 Ok(AgentStatus {
154 command,
155 provider,
156 usage,
157 })
158 }
159
160 async fn check_claude_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
161 let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
162
163 if let Some(reset_time) = usage.next_reset_time() {
164 Ok(AgentLimit::Limited {
165 reset_time: Some(reset_time),
166 })
167 } else {
168 let is_limited = [
169 usage.five_hour.as_ref(),
170 usage.seven_day.as_ref(),
171 usage.seven_day_sonnet.as_ref(),
172 ]
173 .into_iter()
174 .flatten()
175 .any(|w| w.utilization >= 100.0);
176
177 if is_limited {
178 Ok(AgentLimit::Limited { reset_time: None })
179 } else {
180 Ok(AgentLimit::NotLimited)
181 }
182 }
183 }
184
185 async fn check_copilot_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
186 let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
187
188 if quota.is_limited() {
189 Ok(AgentLimit::Limited {
190 reset_time: quota.reset_time,
191 })
192 } else {
193 Ok(AgentLimit::NotLimited)
194 }
195 }
196
197 async fn check_codex_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
198 let usage = crate::codex::CodexClient::fetch_usage(&self.cookies).await?;
199
200 if usage.rate_limit.is_limited() {
201 Ok(AgentLimit::Limited {
202 reset_time: usage.rate_limit.next_reset_time(),
203 })
204 } else {
205 Ok(AgentLimit::NotLimited)
206 }
207 }
208
209 pub fn execute(
210 &self,
211 resolved_args: &[String],
212 extra_args: &[String],
213 ) -> std::io::Result<std::process::ExitStatus> {
214 let mut cmd = std::process::Command::new(self.command());
215 cmd.args(resolved_args);
216 cmd.args(extra_args);
217 if let Some(env) = &self.config.env {
218 cmd.envs(env);
219 }
220 cmd.status()
221 }
222
223 pub fn has_model(&self, model_key: &str) -> bool {
224 match &self.config.models {
225 None => true, Some(m) => m.contains_key(model_key),
227 }
228 }
229
230 pub fn resolved_args(&self, model: Option<&str>) -> Vec<String> {
231 const MODEL_PLACEHOLDER: &str = "{model}";
232 let mut args: Vec<String> = self
233 .config
234 .args
235 .iter()
236 .filter_map(|arg| {
237 if arg.contains(MODEL_PLACEHOLDER) {
238 let model_key = model?;
239 let replacement = self
240 .config
241 .models
242 .as_ref()
243 .and_then(|m| m.get(model_key))
244 .map_or(model_key, |s| s.as_str());
245 Some(arg.replace(MODEL_PLACEHOLDER, replacement))
246 } else {
247 Some(arg.clone())
248 }
249 })
250 .collect();
251
252 if self.config.models.is_none()
254 && let Some(model_key) = model
255 {
256 args.push("--model".to_string());
257 args.push(model_key.to_string());
258 }
259
260 args
261 }
262
263 pub fn mapped_args(&self, args: &[String]) -> Vec<String> {
264 args.iter()
265 .flat_map(|arg| {
266 self.config
267 .arg_maps
268 .get(arg.as_str())
269 .map_or_else(|| std::slice::from_ref(arg), Vec::as_slice)
270 })
271 .cloned()
272 .collect()
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 use std::collections::HashMap;
279
280 use super::*;
281 use crate::codex::{CodexRateLimit, CodexWindow};
282 use crate::config::AgentConfig;
283
284 fn make_agent(
285 models: Option<HashMap<String, String>>,
286 arg_maps: HashMap<String, Vec<String>>,
287 ) -> Agent {
288 Agent::new(
289 AgentConfig {
290 command: "claude".to_string(),
291 args: vec![],
292 models,
293 arg_maps,
294 env: None,
295 provider: None,
296 },
297 vec![],
298 )
299 }
300
301 #[test]
302 fn has_model_returns_true_when_models_is_none() {
303 let agent = make_agent(None, HashMap::new());
304 assert!(agent.has_model("high"));
305 assert!(agent.has_model("anything"));
306 }
307
308 #[test]
309 fn resolved_args_passthrough_when_models_is_none_with_model() {
310 let agent = make_agent(None, HashMap::new());
311 let args = agent.resolved_args(Some("high"));
312 assert_eq!(args, vec!["--model", "high"]);
313 }
314
315 #[test]
316 fn resolved_args_no_model_flag_when_models_is_none_without_model() {
317 let agent = make_agent(None, HashMap::new());
318 let args = agent.resolved_args(None);
319 assert!(!args.contains(&"--model".to_string()));
320 }
321
322 #[test]
323 fn mapped_args_passthrough_when_arg_maps_is_empty() {
324 let agent = make_agent(None, HashMap::new());
325 let args = vec!["--danger".to_string(), "fix bugs".to_string()];
326
327 assert_eq!(agent.mapped_args(&args), args);
328 }
329
330 #[test]
331 fn mapped_args_replaces_matching_tokens() {
332 let mut arg_maps = HashMap::new();
333 arg_maps.insert("--danger".to_string(), vec!["--yolo".to_string()]);
334 let agent = make_agent(None, arg_maps);
335
336 assert_eq!(
337 agent.mapped_args(&["--danger".to_string(), "fix bugs".to_string()]),
338 vec!["--yolo".to_string(), "fix bugs".to_string()]
339 );
340 }
341
342 #[test]
343 fn mapped_args_can_expand_to_multiple_tokens() {
344 let mut arg_maps = HashMap::new();
345 arg_maps.insert(
346 "--danger".to_string(),
347 vec![
348 "--permission-mode".to_string(),
349 "bypassPermissions".to_string(),
350 ],
351 );
352 let agent = make_agent(None, arg_maps);
353
354 assert_eq!(
355 agent.mapped_args(&["--danger".to_string(), "fix bugs".to_string()]),
356 vec![
357 "--permission-mode".to_string(),
358 "bypassPermissions".to_string(),
359 "fix bugs".to_string(),
360 ]
361 );
362 }
363
364 #[test]
365 fn codex_usage_entries_marks_blocking_window_when_only_top_level_limit_is_set() {
366 let limit = CodexRateLimit {
367 allowed: false,
368 limit_reached: false,
369 primary_window: Some(CodexWindow {
370 used_percent: 55.0,
371 limit_window_seconds: 60,
372 reset_after_seconds: 30,
373 reset_at: 100,
374 }),
375 secondary_window: Some(CodexWindow {
376 used_percent: 40.0,
377 limit_window_seconds: 120,
378 reset_after_seconds: 90,
379 reset_at: 200,
380 }),
381 };
382
383 let entries = codex_usage_entries("rate_limit", &limit);
384
385 assert_eq!(entries.len(), 2);
386 assert_eq!(entries[0].entry_type, "rate_limit_primary");
387 assert!(!entries[0].limited);
388 assert_eq!(entries[1].entry_type, "rate_limit_secondary");
389 assert!(entries[1].limited);
390 }
391
392 #[test]
393 fn codex_usage_entries_adds_summary_when_limit_has_no_windows() {
394 let limit = CodexRateLimit {
395 allowed: false,
396 limit_reached: true,
397 primary_window: None,
398 secondary_window: None,
399 };
400
401 let entries = codex_usage_entries("code_review_rate_limit", &limit);
402
403 assert_eq!(entries.len(), 1);
404 assert_eq!(entries[0].entry_type, "code_review_rate_limit");
405 assert!(entries[0].limited);
406 assert_eq!(entries[0].utilization, 100.0);
407 assert_eq!(entries[0].resets_at, None);
408 }
409}