seher/agent/
mod.rs

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    pub domain: String,
10}
11
12#[derive(Debug)]
13pub enum AgentLimit {
14    NotLimited,
15    Limited { reset_time: Option<DateTime<Utc>> },
16}
17
18#[derive(Debug, Serialize)]
19pub struct UsageEntry {
20    #[serde(rename = "type")]
21    pub entry_type: String,
22    pub limited: bool,
23    pub utilization: f64,
24    pub resets_at: Option<DateTime<Utc>>,
25}
26
27#[derive(Debug, Serialize)]
28pub struct AgentStatus {
29    pub command: String,
30    pub usage: Vec<UsageEntry>,
31}
32
33impl Agent {
34    pub fn new(config: AgentConfig, cookies: Vec<Cookie>, domain: String) -> Self {
35        Self {
36            config,
37            cookies,
38            domain,
39        }
40    }
41
42    pub fn command(&self) -> &str {
43        &self.config.command
44    }
45
46    pub fn args(&self) -> &[String] {
47        &self.config.args
48    }
49
50    pub async fn check_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
51        match self.domain.as_str() {
52            "claude.ai" => self.check_claude_limit().await,
53            "github.com" => self.check_copilot_limit().await,
54            _ => Err(format!("Unknown domain: {}", self.domain).into()),
55        }
56    }
57
58    pub async fn fetch_status(&self) -> Result<AgentStatus, Box<dyn std::error::Error>> {
59        match self.domain.as_str() {
60            "claude.ai" => {
61                let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
62                let windows = [
63                    ("five_hour", &usage.five_hour),
64                    ("seven_day", &usage.seven_day),
65                    ("seven_day_sonnet", &usage.seven_day_sonnet),
66                ];
67                let entries = windows
68                    .into_iter()
69                    .filter_map(|(name, w)| {
70                        w.as_ref().map(|w| UsageEntry {
71                            entry_type: name.to_string(),
72                            limited: w.utilization >= 100.0,
73                            utilization: w.utilization,
74                            resets_at: w.resets_at,
75                        })
76                    })
77                    .collect();
78                Ok(AgentStatus {
79                    command: self.config.command.clone(),
80                    usage: entries,
81                })
82            }
83            "github.com" => {
84                let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
85                let entries = vec![
86                    UsageEntry {
87                        entry_type: "chat_utilization".to_string(),
88                        limited: quota.chat_utilization >= 100.0,
89                        utilization: quota.chat_utilization,
90                        resets_at: quota.reset_time,
91                    },
92                    UsageEntry {
93                        entry_type: "premium_utilization".to_string(),
94                        limited: quota.premium_utilization >= 100.0,
95                        utilization: quota.premium_utilization,
96                        resets_at: quota.reset_time,
97                    },
98                ];
99                Ok(AgentStatus {
100                    command: self.config.command.clone(),
101                    usage: entries,
102                })
103            }
104            _ => Err(format!("Unknown domain: {}", self.domain).into()),
105        }
106    }
107
108    async fn check_claude_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
109        let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
110
111        if let Some(reset_time) = usage.next_reset_time() {
112            Ok(AgentLimit::Limited {
113                reset_time: Some(reset_time),
114            })
115        } else {
116            let is_limited = [
117                usage.five_hour.as_ref(),
118                usage.seven_day.as_ref(),
119                usage.seven_day_sonnet.as_ref(),
120            ]
121            .into_iter()
122            .flatten()
123            .any(|w| w.utilization >= 100.0);
124
125            if is_limited {
126                Ok(AgentLimit::Limited { reset_time: None })
127            } else {
128                Ok(AgentLimit::NotLimited)
129            }
130        }
131    }
132
133    async fn check_copilot_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
134        let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
135
136        if quota.is_limited() {
137            Ok(AgentLimit::Limited {
138                reset_time: quota.reset_time,
139            })
140        } else {
141            Ok(AgentLimit::NotLimited)
142        }
143    }
144
145    pub fn execute(&self, extra_args: &[String]) -> std::io::Result<std::process::ExitStatus> {
146        let mut cmd = std::process::Command::new(self.command());
147        cmd.args(self.args());
148        cmd.args(extra_args);
149        cmd.status()
150    }
151}
seher/agent/mod.rs

seher/agent/
mod.rs
