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

seher/agent/
mod.rs
