1use crate::Cookie;
2use crate::config::AgentConfig;
3use chrono::{DateTime, Utc};
4
5pub struct Agent {
6 pub config: AgentConfig,
7 pub cookies: Vec<Cookie>,
8 pub domain: String,
9}
10
11#[derive(Debug)]
12pub enum AgentLimit {
13 NotLimited,
14 Limited { reset_time: Option<DateTime<Utc>> },
15}
16
17impl Agent {
18 pub fn new(config: AgentConfig, cookies: Vec<Cookie>, domain: String) -> Self {
19 Self {
20 config,
21 cookies,
22 domain,
23 }
24 }
25
26 pub fn command(&self) -> &str {
27 &self.config.command
28 }
29
30 pub fn args(&self) -> &[String] {
31 &self.config.args
32 }
33
34 pub async fn check_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
35 match self.domain.as_str() {
36 "claude.ai" => self.check_claude_limit().await,
37 "github.com" => self.check_copilot_limit().await,
38 _ => Err(format!("Unknown domain: {}", self.domain).into()),
39 }
40 }
41
42 async fn check_claude_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
43 let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
44
45 if let Some(reset_time) = usage.next_reset_time() {
46 Ok(AgentLimit::Limited {
47 reset_time: Some(reset_time),
48 })
49 } else {
50 let is_limited = usage
51 .five_hour
52 .as_ref()
53 .map(|w| w.utilization >= 100.0)
54 .unwrap_or(false);
55
56 if is_limited {
57 Ok(AgentLimit::Limited { reset_time: None })
58 } else {
59 Ok(AgentLimit::NotLimited)
60 }
61 }
62 }
63
64 async fn check_copilot_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
65 let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
66
67 if quota.is_limited() {
68 Ok(AgentLimit::Limited {
69 reset_time: quota.reset_time,
70 })
71 } else {
72 Ok(AgentLimit::NotLimited)
73 }
74 }
75
76 pub fn execute(&self, extra_args: &[String]) -> std::io::Result<std::process::ExitStatus> {
77 let mut cmd = std::process::Command::new(self.command());
78 cmd.args(self.args());
79 cmd.args(extra_args);
80 cmd.status()
81 }
82}