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