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}
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(crate::codex::types::CodexWindow::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    #[must_use]
79    pub fn new(config: AgentConfig, cookies: Vec<Cookie>) -> Self {
80        Self { config, cookies }
81    }
82
83    #[must_use]
84    pub fn command(&self) -> &str {
85        &self.config.command
86    }
87
88    #[must_use]
89    pub fn args(&self) -> &[String] {
90        &self.config.args
91    }
92
93    /// # Errors
94    ///
95    /// Returns an error if fetching usage from the provider API fails or the domain is unknown.
96    pub async fn check_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
97        match self.config.resolve_provider() {
98            Some("claude") => self.check_claude_limit().await,
99            Some("codex") => self.check_codex_limit().await,
100            Some("copilot") => self.check_copilot_limit().await,
101            Some("openrouter") => self.check_openrouter_limit().await,
102            Some("glm") => self.check_glm_limit().await,
103            Some("zai") => self.check_zai_limit().await,
104            Some("kimi-k2") => self.check_kimik2_limit().await,
105            Some("warp") => self.check_warp_limit().await,
106            Some("kiro") => self.check_kiro_limit().await,
107            Some("opencode-go") => self.check_opencode_go_limit(),
108            None => Ok(AgentLimit::NotLimited),
109            Some(p) => Err(format!("Unknown provider: {p}").into()),
110        }
111    }
112
113    /// # Errors
114    ///
115    /// Returns an error if fetching usage from the provider API fails or the domain is unknown.
116    #[expect(clippy::too_many_lines)]
117    pub async fn fetch_status(&self) -> Result<AgentStatus, Box<dyn std::error::Error>> {
118        let command = self.config.command.clone();
119        let provider = self.config.resolve_provider().map(ToString::to_string);
120        let usage = match provider.as_deref() {
121            None => vec![],
122            Some("claude") => {
123                let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
124                usage
125                    .all_windows()
126                    .into_iter()
127                    .map(|(name, w)| UsageEntry {
128                        entry_type: name.to_string(),
129                        limited: w.is_limited(),
130                        utilization: w.utilization.unwrap_or(0.0),
131                        resets_at: w.resets_at,
132                    })
133                    .collect()
134            }
135            Some("codex") => {
136                let usage = crate::codex::CodexClient::fetch_usage(&self.cookies).await?;
137                let mut entries = codex_usage_entries("rate_limit", &usage.rate_limit);
138                if let Some(ref cr) = usage.code_review_rate_limit {
139                    entries.extend(codex_usage_entries("code_review_rate_limit", cr));
140                }
141                entries
142            }
143            Some("copilot") => {
144                let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
145                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            }
160            Some("openrouter") => {
161                let management_key = self.openrouter_management_key()?;
162                let credits =
163                    crate::openrouter::OpenRouterClient::fetch_credits(management_key).await?;
164                vec![UsageEntry {
165                    entry_type: "credits".to_string(),
166                    limited: credits.data.is_limited(),
167                    utilization: credits.data.utilization(),
168                    resets_at: None,
169                }]
170            }
171            Some("glm") => {
172                let api_key = self.glm_api_key()?;
173                let quota = crate::glm::GlmClient::fetch_quota(api_key).await?;
174                match quota.data {
175                    Some(data) => data
176                        .limits
177                        .iter()
178                        .map(|l| UsageEntry {
179                            entry_type: l.limit_type.clone(),
180                            limited: l.percentage >= 100,
181                            utilization: f64::from(l.percentage),
182                            resets_at: l.next_reset_time.and_then(DateTime::from_timestamp_millis),
183                        })
184                        .collect(),
185                    None => vec![],
186                }
187            }
188            Some("zai") => {
189                let api_key = self.resolve_env_key("Z_AI_API_KEY")?;
190                let quota_url = self.resolve_optional_env("Z_AI_QUOTA_URL");
191                let quota =
192                    crate::zai::ZaiClient::fetch_quota(&api_key, quota_url.as_deref()).await?;
193                match quota.data {
194                    Some(data) => data
195                        .limits
196                        .iter()
197                        .map(|l| UsageEntry {
198                            entry_type: l.limit_type.clone(),
199                            limited: l.percentage >= 100,
200                            utilization: f64::from(l.percentage),
201                            resets_at: l.next_reset_time.and_then(DateTime::from_timestamp_millis),
202                        })
203                        .collect(),
204                    None => vec![],
205                }
206            }
207            Some("kimi-k2") => {
208                let api_key = self.resolve_env_key("KIMI_K2_API_KEY")?;
209                let credits = crate::kimik2::KimiK2Client::fetch_credits(&api_key).await?;
210                vec![UsageEntry {
211                    entry_type: "credits".to_string(),
212                    limited: credits.is_limited(),
213                    utilization: credits.utilization(),
214                    resets_at: None,
215                }]
216            }
217            Some("warp") => {
218                let api_key = self.resolve_env_key("WARP_API_KEY")?;
219                let info = crate::warp::WarpClient::fetch_limit_info(&api_key).await?;
220                let limit_info = &info.data.get_request_limit_info;
221                vec![UsageEntry {
222                    entry_type: "requests".to_string(),
223                    limited: limit_info.is_limited(),
224                    utilization: limit_info.utilization(),
225                    resets_at: Self::reset_time_from_seconds(limit_info.reset_in_seconds),
226                }]
227            }
228            Some("kiro") => {
229                let info = crate::kiro::KiroClient::fetch_usage().await?;
230                vec![UsageEntry {
231                    entry_type: "requests".to_string(),
232                    limited: info.is_limited(),
233                    utilization: info.utilization(),
234                    resets_at: Self::reset_time_from_seconds(info.reset_in_seconds),
235                }]
236            }
237            Some("opencode-go") => self
238                .opencode_go_usage_snapshot()?
239                .windows
240                .into_iter()
241                .map(|window| UsageEntry {
242                    entry_type: window.entry_type.to_string(),
243                    limited: window.is_limited(),
244                    utilization: window.utilization(),
245                    resets_at: window.resets_at,
246                })
247                .collect(),
248            Some(p) => return Err(format!("Unknown provider: {p}").into()),
249        };
250        Ok(AgentStatus {
251            command,
252            provider,
253            usage,
254        })
255    }
256
257    async fn check_claude_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
258        let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
259        let windows = usage.all_windows();
260
261        let (has_limited, reset_time) =
262            windows
263                .iter()
264                .fold((false, None), |(has_lim, max_t), (_, w)| {
265                    if w.is_limited() {
266                        (true, max_t.max(w.resets_at))
267                    } else {
268                        (has_lim, max_t)
269                    }
270                });
271
272        if has_limited {
273            Ok(AgentLimit::Limited { reset_time })
274        } else {
275            Ok(AgentLimit::NotLimited)
276        }
277    }
278
279    async fn check_copilot_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
280        let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
281
282        if quota.is_limited() {
283            Ok(AgentLimit::Limited {
284                reset_time: quota.reset_time,
285            })
286        } else {
287            Ok(AgentLimit::NotLimited)
288        }
289    }
290
291    fn openrouter_management_key(&self) -> Result<&str, Box<dyn std::error::Error>> {
292        self.config
293            .openrouter_management_key
294            .as_deref()
295            .ok_or_else(|| {
296                "openrouter_management_key is required for OpenRouter provider"
297                    .to_string()
298                    .into()
299            })
300    }
301
302    async fn check_openrouter_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
303        let management_key = self.openrouter_management_key()?;
304        let credits = crate::openrouter::OpenRouterClient::fetch_credits(management_key).await?;
305        if credits.data.is_limited() {
306            Ok(AgentLimit::Limited { reset_time: None })
307        } else {
308            Ok(AgentLimit::NotLimited)
309        }
310    }
311
312    fn glm_api_key(&self) -> Result<&str, Box<dyn std::error::Error>> {
313        self.config.glm_api_key.as_deref().ok_or_else(|| {
314            "glm_api_key is required for GLM provider"
315                .to_string()
316                .into()
317        })
318    }
319
320    async fn check_glm_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
321        let api_key = self.glm_api_key()?;
322        let quota = crate::glm::GlmClient::fetch_quota(api_key).await?;
323        match quota.data {
324            Some(data) if data.is_limited() => {
325                let reset_time = data
326                    .limits
327                    .iter()
328                    .filter_map(|l| l.next_reset_time)
329                    .filter_map(DateTime::from_timestamp_millis)
330                    .max();
331                Ok(AgentLimit::Limited { reset_time })
332            }
333            _ => Ok(AgentLimit::NotLimited),
334        }
335    }
336
337    fn reset_time_from_seconds(secs: Option<i64>) -> Option<DateTime<Utc>> {
338        secs.and_then(|s| Utc::now().checked_add_signed(chrono::Duration::seconds(s)))
339    }
340
341    fn resolve_env_key(&self, key: &str) -> Result<String, Box<dyn std::error::Error>> {
342        // 1. Check agent config env
343        if let Some(env) = &self.config.env
344            && let Some(val) = env.get(key)
345        {
346            return Ok(val.clone());
347        }
348        // 2. Check process environment
349        if let Ok(val) = std::env::var(key) {
350            return Ok(val);
351        }
352        Err(format!("{key} is required for this provider").into())
353    }
354
355    fn resolve_optional_env(&self, key: &str) -> Option<String> {
356        self.config
357            .env
358            .as_ref()
359            .and_then(|env| env.get(key).cloned())
360            .or_else(|| std::env::var(key).ok())
361    }
362
363    async fn check_zai_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
364        let api_key = self.resolve_env_key("Z_AI_API_KEY")?;
365        let quota_url = self.resolve_optional_env("Z_AI_QUOTA_URL");
366        let quota = crate::zai::ZaiClient::fetch_quota(&api_key, quota_url.as_deref()).await?;
367        match quota.data {
368            Some(data) if data.is_limited() => {
369                let reset_time = data
370                    .limits
371                    .iter()
372                    .filter_map(|l| l.next_reset_time)
373                    .filter_map(DateTime::from_timestamp_millis)
374                    .max();
375                Ok(AgentLimit::Limited { reset_time })
376            }
377            _ => Ok(AgentLimit::NotLimited),
378        }
379    }
380
381    async fn check_kimik2_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
382        let api_key = self.resolve_env_key("KIMI_K2_API_KEY")?;
383        let credits = crate::kimik2::KimiK2Client::fetch_credits(&api_key).await?;
384        if credits.is_limited() {
385            Ok(AgentLimit::Limited { reset_time: None })
386        } else {
387            Ok(AgentLimit::NotLimited)
388        }
389    }
390
391    async fn check_warp_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
392        let api_key = self.resolve_env_key("WARP_API_KEY")?;
393        let info = crate::warp::WarpClient::fetch_limit_info(&api_key).await?;
394        let limit_info = &info.data.get_request_limit_info;
395        if limit_info.is_limited() {
396            Ok(AgentLimit::Limited {
397                reset_time: Self::reset_time_from_seconds(limit_info.reset_in_seconds),
398            })
399        } else {
400            Ok(AgentLimit::NotLimited)
401        }
402    }
403
404    async fn check_kiro_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
405        let info = crate::kiro::KiroClient::fetch_usage().await?;
406        if info.is_limited() {
407            Ok(AgentLimit::Limited {
408                reset_time: Self::reset_time_from_seconds(info.reset_in_seconds),
409            })
410        } else {
411            Ok(AgentLimit::NotLimited)
412        }
413    }
414
415    fn check_opencode_go_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
416        let snapshot = self.opencode_go_usage_snapshot()?;
417        if snapshot
418            .windows
419            .iter()
420            .any(crate::opencode_go::OpencodeGoUsageWindow::is_limited)
421        {
422            Ok(AgentLimit::Limited {
423                reset_time: snapshot.reset_time(),
424            })
425        } else {
426            Ok(AgentLimit::NotLimited)
427        }
428    }
429
430    async fn check_codex_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
431        let usage = crate::codex::CodexClient::fetch_usage(&self.cookies).await?;
432
433        if usage.rate_limit.is_limited() {
434            Ok(AgentLimit::Limited {
435                reset_time: usage.rate_limit.next_reset_time(),
436            })
437        } else {
438            Ok(AgentLimit::NotLimited)
439        }
440    }
441
442    /// # Errors
443    ///
444    /// Returns an error if spawning or waiting on the child process fails.
445    pub fn execute(
446        &self,
447        resolved_args: &[String],
448        extra_args: &[String],
449    ) -> std::io::Result<std::process::ExitStatus> {
450        if let Some((cmd, args)) = self.config.pre_command.split_first() {
451            let mut pre_cmd = std::process::Command::new(cmd);
452            pre_cmd.args(args);
453            if let Some(env) = &self.config.env {
454                pre_cmd.envs(env);
455            }
456            let status = pre_cmd.status()?;
457            if !status.success() {
458                return Ok(status);
459            }
460        }
461        let mut cmd = std::process::Command::new(self.command());
462        cmd.args(resolved_args);
463        cmd.args(extra_args);
464        if let Some(env) = &self.config.env {
465            cmd.envs(env);
466        }
467        cmd.status()
468    }
469
470    #[must_use]
471    pub fn has_model(&self, model_key: &str) -> bool {
472        self.config.has_model(model_key)
473    }
474
475    #[must_use]
476    pub fn resolved_args(&self, model: Option<&str>) -> Vec<String> {
477        const MODEL_PLACEHOLDER: &str = "{model}";
478        let mut args: Vec<String> = self
479            .config
480            .args
481            .iter()
482            .filter_map(|arg| {
483                if arg.contains(MODEL_PLACEHOLDER) {
484                    let model_key = model?;
485                    let replacement = self
486                        .config
487                        .models
488                        .as_ref()
489                        .and_then(|m| m.get(model_key))
490                        .map_or(model_key, |s| s.as_str());
491                    Some(arg.replace(MODEL_PLACEHOLDER, replacement))
492                } else {
493                    Some(arg.clone())
494                }
495            })
496            .collect();
497
498        // If models map is not set, pass --model <value> through as-is
499        if self.config.models.is_none()
500            && let Some(model_key) = model
501        {
502            args.push("--model".to_string());
503            args.push(model_key.to_string());
504        }
505
506        args
507    }
508
509    #[must_use]
510    pub fn mapped_args(&self, args: &[String]) -> Vec<String> {
511        args.iter()
512            .flat_map(|arg| {
513                self.config
514                    .arg_maps
515                    .get(arg.as_str())
516                    .map_or_else(|| std::slice::from_ref(arg), Vec::as_slice)
517            })
518            .cloned()
519            .collect()
520    }
521
522    fn opencode_go_usage_snapshot(
523        &self,
524    ) -> Result<crate::opencode_go::OpencodeGoUsageSnapshot, Box<dyn std::error::Error>> {
525        let db_path = self.resolve_optional_env("SEHER_OPENCODE_DB_PATH");
526        let auth_path = self.resolve_optional_env("SEHER_OPENCODE_AUTH_PATH");
527        Ok(
528            crate::opencode_go::OpencodeGoUsageStore::fetch_usage_with_paths_at(
529                db_path.as_deref().map(std::path::Path::new),
530                auth_path.as_deref().map(std::path::Path::new),
531                Utc::now(),
532            )?,
533        )
534    }
535}
536
537#[cfg(test)]
538mod tests {
539    use std::collections::HashMap;
540
541    use super::*;
542    use crate::codex::{CodexRateLimit, CodexWindow};
543    use crate::config::AgentConfig;
544
545    fn make_agent(
546        models: Option<HashMap<String, String>>,
547        arg_maps: HashMap<String, Vec<String>>,
548    ) -> Agent {
549        Agent::new(
550            AgentConfig {
551                command: "claude".to_string(),
552                args: vec![],
553                models,
554                arg_maps,
555                env: None,
556                provider: None,
557                openrouter_management_key: None,
558                glm_api_key: None,
559                pre_command: vec![],
560            },
561            vec![],
562        )
563    }
564
565    #[test]
566    fn has_model_returns_true_when_models_is_none() {
567        let agent = make_agent(None, HashMap::new());
568        assert!(agent.has_model("high"));
569        assert!(agent.has_model("anything"));
570    }
571
572    #[test]
573    fn resolved_args_passthrough_when_models_is_none_with_model() {
574        let agent = make_agent(None, HashMap::new());
575        let args = agent.resolved_args(Some("high"));
576        assert_eq!(args, vec!["--model", "high"]);
577    }
578
579    #[test]
580    fn resolved_args_no_model_flag_when_models_is_none_without_model() {
581        let agent = make_agent(None, HashMap::new());
582        let args = agent.resolved_args(None);
583        assert!(!args.contains(&"--model".to_string()));
584    }
585
586    #[test]
587    fn mapped_args_passthrough_when_arg_maps_is_empty() {
588        let agent = make_agent(None, HashMap::new());
589        let args = vec!["--danger".to_string(), "fix bugs".to_string()];
590
591        assert_eq!(agent.mapped_args(&args), args);
592    }
593
594    #[test]
595    fn mapped_args_replaces_matching_tokens() {
596        let mut arg_maps = HashMap::new();
597        arg_maps.insert("--danger".to_string(), vec!["--yolo".to_string()]);
598        let agent = make_agent(None, arg_maps);
599
600        assert_eq!(
601            agent.mapped_args(&["--danger".to_string(), "fix bugs".to_string()]),
602            vec!["--yolo".to_string(), "fix bugs".to_string()]
603        );
604    }
605
606    #[test]
607    fn mapped_args_can_expand_to_multiple_tokens() {
608        let mut arg_maps = HashMap::new();
609        arg_maps.insert(
610            "--danger".to_string(),
611            vec![
612                "--permission-mode".to_string(),
613                "bypassPermissions".to_string(),
614            ],
615        );
616        let agent = make_agent(None, arg_maps);
617
618        assert_eq!(
619            agent.mapped_args(&["--danger".to_string(), "fix bugs".to_string()]),
620            vec![
621                "--permission-mode".to_string(),
622                "bypassPermissions".to_string(),
623                "fix bugs".to_string(),
624            ]
625        );
626    }
627
628    #[test]
629    fn codex_usage_entries_marks_blocking_window_when_only_top_level_limit_is_set() {
630        let limit = CodexRateLimit {
631            allowed: false,
632            limit_reached: false,
633            primary_window: Some(CodexWindow {
634                used_percent: 55.0,
635                limit_window_seconds: 60,
636                reset_after_seconds: 30,
637                reset_at: 100,
638            }),
639            secondary_window: Some(CodexWindow {
640                used_percent: 40.0,
641                limit_window_seconds: 120,
642                reset_after_seconds: 90,
643                reset_at: 200,
644            }),
645        };
646
647        let entries = codex_usage_entries("rate_limit", &limit);
648
649        assert_eq!(entries.len(), 2);
650        assert_eq!(entries[0].entry_type, "rate_limit_primary");
651        assert!(!entries[0].limited);
652        assert_eq!(entries[1].entry_type, "rate_limit_secondary");
653        assert!(entries[1].limited);
654    }
655
656    #[test]
657    fn codex_usage_entries_adds_summary_when_limit_has_no_windows() {
658        let limit = CodexRateLimit {
659            allowed: false,
660            limit_reached: true,
661            primary_window: None,
662            secondary_window: None,
663        };
664
665        let entries = codex_usage_entries("code_review_rate_limit", &limit);
666
667        assert_eq!(entries.len(), 1);
668        assert_eq!(entries[0].entry_type, "code_review_rate_limit");
669        assert!(entries[0].limited);
670        assert!((entries[0].utilization - 100.0).abs() < f64::EPSILON);
671        assert_eq!(entries[0].resets_at, None);
672    }
673
674    // -----------------------------------------------------------------------
675    // OpenRouter dispatch tests
676    // These tests verify that check_limit() / fetch_status() correctly route
677    // to the openrouter handler when provider == "openrouter", and that a
678    // missing management key causes an immediate error (no HTTP call made).
679    // -----------------------------------------------------------------------
680
681    fn make_openrouter_agent(management_key: Option<&str>) -> Agent {
682        Agent::new(
683            AgentConfig {
684                command: "myai".to_string(),
685                args: vec![],
686                models: None,
687                arg_maps: HashMap::new(),
688                env: None,
689                provider: Some(crate::config::ProviderConfig::Explicit(
690                    "openrouter".to_string(),
691                )),
692                openrouter_management_key: management_key.map(str::to_string),
693                glm_api_key: None,
694                pre_command: vec![],
695            },
696            vec![],
697        )
698    }
699
700    fn make_agent_with_pre_command(pre_command: Vec<String>, main_command: &str) -> Agent {
701        Agent::new(
702            AgentConfig {
703                command: main_command.to_string(),
704                args: vec![],
705                models: None,
706                arg_maps: HashMap::new(),
707                env: None,
708                provider: None,
709                openrouter_management_key: None,
710                glm_api_key: None,
711                pre_command,
712            },
713            vec![],
714        )
715    }
716
717    #[test]
718    #[cfg(unix)]
719    fn execute_runs_main_command_when_pre_command_succeeds() -> TestResult {
720        // pre_command: true (always exits 0), main: true
721        let agent = make_agent_with_pre_command(vec!["true".to_string()], "true");
722        let status = agent.execute(&[], &[])?;
723        assert!(status.success());
724        Ok(())
725    }
726
727    #[test]
728    #[cfg(unix)]
729    fn execute_skips_main_command_when_pre_command_fails() -> TestResult {
730        // pre_command: false (always exits non-0), main: true
731        let agent = make_agent_with_pre_command(vec!["false".to_string()], "true");
732        let status = agent.execute(&[], &[])?;
733        assert!(!status.success());
734        Ok(())
735    }
736
737    type TestResult = Result<(), Box<dyn std::error::Error>>;
738
739    #[tokio::test(flavor = "current_thread")]
740    async fn check_limit_openrouter_returns_error_when_management_key_is_missing() -> TestResult {
741        // Given: openrouter agent with no management key configured
742        let agent = make_openrouter_agent(None);
743
744        // When: check_limit is called
745        let result = agent.check_limit().await;
746
747        // Then: error mentions the missing key -- no HTTP call should be made
748        let err_msg = result.err().ok_or("expected Err")?.to_string();
749        assert!(err_msg.contains("openrouter_management_key"));
750        Ok(())
751    }
752
753    #[tokio::test(flavor = "current_thread")]
754    async fn fetch_status_openrouter_returns_error_when_management_key_is_missing() -> TestResult {
755        // Given: openrouter agent with no management key configured
756        let agent = make_openrouter_agent(None);
757
758        // When: fetch_status is called
759        let result = agent.fetch_status().await;
760
761        // Then: error mentions the missing key -- no HTTP call should be made
762        let err_msg = result.err().ok_or("expected Err")?.to_string();
763        assert!(err_msg.contains("openrouter_management_key"));
764        Ok(())
765    }
766
767    fn make_api_key_agent(provider: &str) -> Agent {
768        Agent::new(
769            AgentConfig {
770                command: "myai".to_string(),
771                args: vec![],
772                models: None,
773                arg_maps: HashMap::new(),
774                env: None,
775                provider: Some(crate::config::ProviderConfig::Explicit(
776                    provider.to_string(),
777                )),
778                openrouter_management_key: None,
779                glm_api_key: None,
780                pre_command: vec![],
781            },
782            vec![],
783        )
784    }
785
786    // -- zai --
787
788    #[tokio::test(flavor = "current_thread")]
789    async fn check_limit_zai_returns_error_when_api_key_is_missing() -> TestResult {
790        let agent = make_api_key_agent("zai");
791        let result = agent.check_limit().await;
792        let err_msg = result.err().ok_or("expected Err")?.to_string();
793        assert!(
794            err_msg.contains("Z_AI_API_KEY"),
795            "error should mention Z_AI_API_KEY, got: {err_msg}"
796        );
797        Ok(())
798    }
799
800    #[tokio::test(flavor = "current_thread")]
801    async fn fetch_status_zai_returns_error_when_api_key_is_missing() -> TestResult {
802        let agent = make_api_key_agent("zai");
803        let result = agent.fetch_status().await;
804        let err_msg = result.err().ok_or("expected Err")?.to_string();
805        assert!(
806            err_msg.contains("Z_AI_API_KEY"),
807            "error should mention Z_AI_API_KEY, got: {err_msg}"
808        );
809        Ok(())
810    }
811
812    // -- kimi-k2 --
813
814    #[tokio::test(flavor = "current_thread")]
815    async fn check_limit_kimik2_returns_error_when_api_key_is_missing() -> TestResult {
816        let agent = make_api_key_agent("kimi-k2");
817        let result = agent.check_limit().await;
818        let err_msg = result.err().ok_or("expected Err")?.to_string();
819        assert!(
820            err_msg.contains("KIMI_K2_API_KEY"),
821            "error should mention KIMI_K2_API_KEY, got: {err_msg}"
822        );
823        Ok(())
824    }
825
826    #[tokio::test(flavor = "current_thread")]
827    async fn fetch_status_kimik2_returns_error_when_api_key_is_missing() -> TestResult {
828        let agent = make_api_key_agent("kimi-k2");
829        let result = agent.fetch_status().await;
830        let err_msg = result.err().ok_or("expected Err")?.to_string();
831        assert!(
832            err_msg.contains("KIMI_K2_API_KEY"),
833            "error should mention KIMI_K2_API_KEY, got: {err_msg}"
834        );
835        Ok(())
836    }
837
838    // -- warp --
839
840    #[tokio::test(flavor = "current_thread")]
841    async fn check_limit_warp_returns_error_when_api_key_is_missing() -> TestResult {
842        let agent = make_api_key_agent("warp");
843        let result = agent.check_limit().await;
844        let err_msg = result.err().ok_or("expected Err")?.to_string();
845        assert!(
846            err_msg.contains("WARP_API_KEY"),
847            "error should mention WARP_API_KEY, got: {err_msg}"
848        );
849        Ok(())
850    }
851
852    #[tokio::test(flavor = "current_thread")]
853    async fn fetch_status_warp_returns_error_when_api_key_is_missing() -> TestResult {
854        let agent = make_api_key_agent("warp");
855        let result = agent.fetch_status().await;
856        let err_msg = result.err().ok_or("expected Err")?.to_string();
857        assert!(
858            err_msg.contains("WARP_API_KEY"),
859            "error should mention WARP_API_KEY, got: {err_msg}"
860        );
861        Ok(())
862    }
863
864    // -- kiro (CLI-based, no API key needed for dispatch, but must not panic) --
865
866    #[tokio::test(flavor = "current_thread")]
867    async fn check_limit_kiro_returns_error_when_command_not_found() -> TestResult {
868        let agent = make_api_key_agent("kiro");
869        let result = agent.check_limit().await;
870        assert!(result.is_err(), "kiro without CLI should return an error");
871        Ok(())
872    }
873
874    #[tokio::test(flavor = "current_thread")]
875    async fn fetch_status_kiro_returns_error_when_command_not_found() -> TestResult {
876        let agent = make_api_key_agent("kiro");
877        let result = agent.fetch_status().await;
878        assert!(result.is_err(), "kiro without CLI should return an error");
879        Ok(())
880    }
881
882    #[tokio::test(flavor = "current_thread")]
883    async fn check_limit_opencode_go_uses_local_history() -> TestResult {
884        let tmp = tempfile::tempdir()?;
885        let db_path = tmp.path().join("opencode.db");
886        let conn = rusqlite::Connection::open(&db_path)?;
887        conn.execute("CREATE TABLE message (data TEXT NOT NULL)", [])?;
888        conn.execute(
889            "INSERT INTO message (data) VALUES (?1)",
890            [r#"{"role":"assistant","providerID":"opencode-go","cost":6.5,"time":{"completed":4102448400000}}"#],
891        )?;
892        conn.execute(
893            "INSERT INTO message (data) VALUES (?1)",
894            [r#"{"role":"assistant","providerID":"opencode-go","cost":6.0,"time":{"completed":4102461000000}}"#],
895        )?;
896        drop(conn);
897
898        let mut agent = make_api_key_agent("opencode-go");
899        agent.config.env = Some(HashMap::from([(
900            "SEHER_OPENCODE_DB_PATH".to_string(),
901            db_path.display().to_string(),
902        )]));
903        let result = agent.check_limit().await?;
904        assert!(matches!(result, AgentLimit::Limited { .. }));
905        Ok(())
906    }
907
908    #[tokio::test(flavor = "current_thread")]
909    async fn fetch_status_opencode_go_returns_usage_windows() -> TestResult {
910        let tmp = tempfile::tempdir()?;
911        let db_path = tmp.path().join("opencode.db");
912        let conn = rusqlite::Connection::open(&db_path)?;
913        conn.execute("CREATE TABLE message (data TEXT NOT NULL)", [])?;
914        conn.execute(
915            "INSERT INTO message (data) VALUES (?1)",
916            [r#"{"role":"assistant","providerID":"opencode-go","cost":2.25,"time":{"completed":4102461000000}}"#],
917        )?;
918        drop(conn);
919
920        let mut agent = make_api_key_agent("opencode-go");
921        agent.config.env = Some(HashMap::from([(
922            "SEHER_OPENCODE_DB_PATH".to_string(),
923            db_path.display().to_string(),
924        )]));
925        let status = agent.fetch_status().await?;
926        assert_eq!(status.provider.as_deref(), Some("opencode-go"));
927        assert_eq!(status.usage.len(), 3);
928        assert!(
929            status
930                .usage
931                .iter()
932                .any(|entry| entry.entry_type == "five_hour_spend")
933        );
934        assert!(
935            status
936                .usage
937                .iter()
938                .any(|entry| entry.entry_type == "weekly_spend")
939        );
940        assert!(
941            status
942                .usage
943                .iter()
944                .any(|entry| entry.entry_type == "monthly_spend")
945        );
946        Ok(())
947    }
948
949    // -- unknown provider still errors --
950
951    #[tokio::test(flavor = "current_thread")]
952    async fn check_limit_unknown_provider_returns_error() -> TestResult {
953        let agent = make_api_key_agent("nonexistent-provider");
954        let result = agent.check_limit().await;
955        let err_msg = result.err().ok_or("expected Err")?.to_string();
956        assert!(err_msg.contains("Unknown provider"), "got: {err_msg}");
957        Ok(())
958    }
959}
seher/agent/mod.rs

seher/agent/
mod.rs
