Skip to main content

seher/agent/
mod.rs

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}