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, Clone)]
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    /// # Errors
89    ///
90    /// Returns an error if fetching usage from the provider API fails or the domain is unknown.
91    pub async fn check_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
92        match self.config.resolve_provider() {
93            Some("claude") => self.check_claude_limit().await,
94            Some("codex") => self.check_codex_limit().await,
95            Some("copilot") => self.check_copilot_limit().await,
96            Some("openrouter") => self.check_openrouter_limit().await,
97            Some("glm") => self.check_glm_limit().await,
98            Some("zai") => self.check_zai_limit().await,
99            Some("kimi-k2") => self.check_kimik2_limit().await,
100            Some("warp") => self.check_warp_limit().await,
101            Some("kiro") => self.check_kiro_limit().await,
102            Some("opencode-go") => self.check_opencode_go_limit().await,
103            None => Ok(AgentLimit::NotLimited),
104            Some(p) => Err(format!("Unknown provider: {p}").into()),
105        }
106    }
107
108    /// # Errors
109    ///
110    /// Returns an error if fetching usage from the provider API fails or the domain is unknown.
111    #[expect(clippy::too_many_lines)]
112    pub async fn fetch_status(&self) -> Result<AgentStatus, Box<dyn std::error::Error>> {
113        let command = self.config.command.clone();
114        let provider = self.config.resolve_provider().map(ToString::to_string);
115        let usage = match provider.as_deref() {
116            None => vec![],
117            Some("claude") => {
118                let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
119                usage
120                    .all_windows()
121                    .into_iter()
122                    .map(|(name, w)| UsageEntry {
123                        entry_type: name.to_string(),
124                        limited: w.is_limited(),
125                        utilization: w.utilization.unwrap_or(0.0),
126                        resets_at: w.resets_at,
127                    })
128                    .collect()
129            }
130            Some("codex") => {
131                let usage = crate::codex::CodexClient::fetch_usage(&self.cookies).await?;
132                let mut entries = codex_usage_entries("rate_limit", &usage.rate_limit);
133                if let Some(ref cr) = usage.code_review_rate_limit {
134                    entries.extend(codex_usage_entries("code_review_rate_limit", cr));
135                }
136                entries
137            }
138            Some("copilot") => {
139                let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
140                vec![
141                    UsageEntry {
142                        entry_type: "chat_utilization".to_string(),
143                        limited: quota.chat_utilization >= 100.0,
144                        utilization: quota.chat_utilization,
145                        resets_at: quota.reset_time,
146                    },
147                    UsageEntry {
148                        entry_type: "premium_utilization".to_string(),
149                        limited: quota.premium_utilization >= 100.0,
150                        utilization: quota.premium_utilization,
151                        resets_at: quota.reset_time,
152                    },
153                ]
154            }
155            Some("openrouter") => {
156                let management_key = self.openrouter_management_key()?;
157                let credits =
158                    crate::openrouter::OpenRouterClient::fetch_credits(management_key).await?;
159                vec![UsageEntry {
160                    entry_type: "credits".to_string(),
161                    limited: credits.data.is_limited(),
162                    utilization: credits.data.utilization(),
163                    resets_at: None,
164                }]
165            }
166            Some("glm") => {
167                let api_key = self.glm_api_key()?;
168                let quota = crate::glm::GlmClient::fetch_quota(api_key).await?;
169                match quota.data {
170                    Some(data) => data
171                        .limits
172                        .iter()
173                        .map(|l| UsageEntry {
174                            entry_type: l.limit_type.clone(),
175                            limited: l.percentage >= 100,
176                            utilization: f64::from(l.percentage),
177                            resets_at: l.next_reset_time.and_then(DateTime::from_timestamp_millis),
178                        })
179                        .collect(),
180                    None => vec![],
181                }
182            }
183            Some("zai") => {
184                let api_key = self.resolve_env_key("Z_AI_API_KEY")?;
185                let quota_url = self.resolve_optional_env("Z_AI_QUOTA_URL");
186                let quota =
187                    crate::zai::ZaiClient::fetch_quota(&api_key, quota_url.as_deref()).await?;
188                match quota.data {
189                    Some(data) => data
190                        .limits
191                        .iter()
192                        .map(|l| UsageEntry {
193                            entry_type: l.limit_type.clone(),
194                            limited: l.percentage >= 100,
195                            utilization: f64::from(l.percentage),
196                            resets_at: l.next_reset_time.and_then(DateTime::from_timestamp_millis),
197                        })
198                        .collect(),
199                    None => vec![],
200                }
201            }
202            Some("kimi-k2") => {
203                let api_key = self.resolve_env_key("KIMI_K2_API_KEY")?;
204                let credits = crate::kimik2::KimiK2Client::fetch_credits(&api_key).await?;
205                vec![UsageEntry {
206                    entry_type: "credits".to_string(),
207                    limited: credits.is_limited(),
208                    utilization: credits.utilization(),
209                    resets_at: None,
210                }]
211            }
212            Some("warp") => {
213                let api_key = self.resolve_env_key("WARP_API_KEY")?;
214                let info = crate::warp::WarpClient::fetch_limit_info(&api_key).await?;
215                let limit_info = &info.data.get_request_limit_info;
216                vec![UsageEntry {
217                    entry_type: "requests".to_string(),
218                    limited: limit_info.is_limited(),
219                    utilization: limit_info.utilization(),
220                    resets_at: Self::reset_time_from_seconds(limit_info.reset_in_seconds),
221                }]
222            }
223            Some("kiro") => {
224                let info = crate::kiro::KiroClient::fetch_usage().await?;
225                vec![UsageEntry {
226                    entry_type: "requests".to_string(),
227                    limited: info.is_limited(),
228                    utilization: info.utilization(),
229                    resets_at: Self::reset_time_from_seconds(info.reset_in_seconds),
230                }]
231            }
232            Some("opencode-go") => self
233                .opencode_go_usage_snapshot()?
234                .windows
235                .into_iter()
236                .map(|window| UsageEntry {
237                    entry_type: window.entry_type.to_string(),
238                    limited: window.is_limited(),
239                    utilization: window.utilization(),
240                    resets_at: window.resets_at,
241                })
242                .collect(),
243            Some(p) => return Err(format!("Unknown provider: {p}").into()),
244        };
245        Ok(AgentStatus {
246            command,
247            provider,
248            usage,
249        })
250    }
251
252    async fn check_claude_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
253        let usage = crate::claude::ClaudeClient::fetch_usage(&self.cookies).await?;
254        let windows = usage.all_windows();
255
256        let (has_limited, reset_time) =
257            windows
258                .iter()
259                .fold((false, None), |(has_lim, max_t), (_, w)| {
260                    if w.is_limited() {
261                        (true, max_t.max(w.resets_at))
262                    } else {
263                        (has_lim, max_t)
264                    }
265                });
266
267        if has_limited {
268            Ok(AgentLimit::Limited { reset_time })
269        } else {
270            Ok(AgentLimit::NotLimited)
271        }
272    }
273
274    async fn check_copilot_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
275        let quota = crate::copilot::CopilotClient::fetch_quota(&self.cookies).await?;
276
277        if quota.is_limited() {
278            Ok(AgentLimit::Limited {
279                reset_time: quota.reset_time,
280            })
281        } else {
282            Ok(AgentLimit::NotLimited)
283        }
284    }
285
286    fn openrouter_management_key(&self) -> Result<&str, Box<dyn std::error::Error>> {
287        self.config
288            .openrouter_management_key
289            .as_deref()
290            .ok_or_else(|| {
291                "openrouter_management_key is required for OpenRouter provider"
292                    .to_string()
293                    .into()
294            })
295    }
296
297    async fn check_openrouter_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
298        let management_key = self.openrouter_management_key()?;
299        let credits = crate::openrouter::OpenRouterClient::fetch_credits(management_key).await?;
300        if credits.data.is_limited() {
301            Ok(AgentLimit::Limited { reset_time: None })
302        } else {
303            Ok(AgentLimit::NotLimited)
304        }
305    }
306
307    fn glm_api_key(&self) -> Result<&str, Box<dyn std::error::Error>> {
308        self.config.glm_api_key.as_deref().ok_or_else(|| {
309            "glm_api_key is required for GLM provider"
310                .to_string()
311                .into()
312        })
313    }
314
315    async fn check_glm_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
316        let api_key = self.glm_api_key()?;
317        let quota = crate::glm::GlmClient::fetch_quota(api_key).await?;
318        match quota.data {
319            Some(data) if data.is_limited() => {
320                let reset_time = data
321                    .limits
322                    .iter()
323                    .filter_map(|l| l.next_reset_time)
324                    .filter_map(DateTime::from_timestamp_millis)
325                    .max();
326                Ok(AgentLimit::Limited { reset_time })
327            }
328            _ => Ok(AgentLimit::NotLimited),
329        }
330    }
331
332    fn reset_time_from_seconds(secs: Option<i64>) -> Option<DateTime<Utc>> {
333        secs.and_then(|s| Utc::now().checked_add_signed(chrono::Duration::seconds(s)))
334    }
335
336    fn resolve_env_key(&self, key: &str) -> Result<String, Box<dyn std::error::Error>> {
337        // 1. Check agent config env
338        if let Some(env) = &self.config.env
339            && let Some(val) = env.get(key)
340        {
341            return Ok(val.clone());
342        }
343        // 2. Check process environment
344        if let Ok(val) = std::env::var(key) {
345            return Ok(val);
346        }
347        Err(format!("{key} is required for this provider").into())
348    }
349
350    fn resolve_optional_env(&self, key: &str) -> Option<String> {
351        self.config
352            .env
353            .as_ref()
354            .and_then(|env| env.get(key).cloned())
355            .or_else(|| std::env::var(key).ok())
356    }
357
358    async fn check_zai_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
359        let api_key = self.resolve_env_key("Z_AI_API_KEY")?;
360        let quota_url = self.resolve_optional_env("Z_AI_QUOTA_URL");
361        let quota = crate::zai::ZaiClient::fetch_quota(&api_key, quota_url.as_deref()).await?;
362        match quota.data {
363            Some(data) if data.is_limited() => {
364                let reset_time = data
365                    .limits
366                    .iter()
367                    .filter_map(|l| l.next_reset_time)
368                    .filter_map(DateTime::from_timestamp_millis)
369                    .max();
370                Ok(AgentLimit::Limited { reset_time })
371            }
372            _ => Ok(AgentLimit::NotLimited),
373        }
374    }
375
376    async fn check_kimik2_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
377        let api_key = self.resolve_env_key("KIMI_K2_API_KEY")?;
378        let credits = crate::kimik2::KimiK2Client::fetch_credits(&api_key).await?;
379        if credits.is_limited() {
380            Ok(AgentLimit::Limited { reset_time: None })
381        } else {
382            Ok(AgentLimit::NotLimited)
383        }
384    }
385
386    async fn check_warp_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
387        let api_key = self.resolve_env_key("WARP_API_KEY")?;
388        let info = crate::warp::WarpClient::fetch_limit_info(&api_key).await?;
389        let limit_info = &info.data.get_request_limit_info;
390        if limit_info.is_limited() {
391            Ok(AgentLimit::Limited {
392                reset_time: Self::reset_time_from_seconds(limit_info.reset_in_seconds),
393            })
394        } else {
395            Ok(AgentLimit::NotLimited)
396        }
397    }
398
399    async fn check_kiro_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
400        let info = crate::kiro::KiroClient::fetch_usage().await?;
401        if info.is_limited() {
402            Ok(AgentLimit::Limited {
403                reset_time: Self::reset_time_from_seconds(info.reset_in_seconds),
404            })
405        } else {
406            Ok(AgentLimit::NotLimited)
407        }
408    }
409
410    async fn check_opencode_go_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
411        // Prefer the authoritative hosted dashboard (real account usage) when an
412        // opencode.ai session cookie is available; the local SQLite heuristic
413        // only sees this machine's spend and badly undercounts a multi-device
414        // account. Fall back to the local snapshot if the remote fetch fails
415        // (no cookie, signed out, network error).
416        // Fall through to the local heuristic if the remote fetch fails
417        // (no cookie, signed out, network error).
418        if !self.cookies.is_empty()
419            && let Ok(usage) = crate::opencode_go::fetch_remote_usage(&self.cookies).await
420        {
421            return Ok(if usage.is_limited() {
422                AgentLimit::Limited {
423                    reset_time: usage.reset_time(Utc::now()),
424                }
425            } else {
426                AgentLimit::NotLimited
427            });
428        }
429
430        let snapshot = self.opencode_go_usage_snapshot()?;
431        if snapshot
432            .windows
433            .iter()
434            .any(crate::opencode_go::OpencodeGoUsageWindow::is_limited)
435        {
436            Ok(AgentLimit::Limited {
437                reset_time: snapshot.reset_time(),
438            })
439        } else {
440            Ok(AgentLimit::NotLimited)
441        }
442    }
443
444    async fn check_codex_limit(&self) -> Result<AgentLimit, Box<dyn std::error::Error>> {
445        let usage = crate::codex::CodexClient::fetch_usage(&self.cookies).await?;
446
447        if usage.rate_limit.is_limited() {
448            Ok(AgentLimit::Limited {
449                reset_time: usage.rate_limit.next_reset_time(),
450            })
451        } else {
452            Ok(AgentLimit::NotLimited)
453        }
454    }
455
456    fn opencode_go_usage_snapshot(
457        &self,
458    ) -> Result<crate::opencode_go::OpencodeGoUsageSnapshot, Box<dyn std::error::Error>> {
459        let db_path = self.resolve_optional_env("SEHER_OPENCODE_DB_PATH");
460        let auth_path = self.resolve_optional_env("SEHER_OPENCODE_AUTH_PATH");
461        Ok(
462            crate::opencode_go::OpencodeGoUsageStore::fetch_usage_with_paths_at(
463                db_path.as_deref().map(std::path::Path::new),
464                auth_path.as_deref().map(std::path::Path::new),
465                Utc::now(),
466            )?,
467        )
468    }
469}
470
471#[cfg(test)]
472mod tests {
473    use std::collections::HashMap;
474
475    use super::*;
476    use crate::codex::{CodexRateLimit, CodexWindow};
477    use crate::config::AgentConfig;
478
479    #[test]
480    fn codex_usage_entries_marks_blocking_window_when_only_top_level_limit_is_set() {
481        let limit = CodexRateLimit {
482            allowed: false,
483            limit_reached: false,
484            primary_window: Some(CodexWindow {
485                used_percent: 55.0,
486                limit_window_seconds: 60,
487                reset_after_seconds: 30,
488                reset_at: 100,
489            }),
490            secondary_window: Some(CodexWindow {
491                used_percent: 40.0,
492                limit_window_seconds: 120,
493                reset_after_seconds: 90,
494                reset_at: 200,
495            }),
496        };
497
498        let entries = codex_usage_entries("rate_limit", &limit);
499
500        assert_eq!(entries.len(), 2);
501        assert_eq!(entries[0].entry_type, "rate_limit_primary");
502        assert!(!entries[0].limited);
503        assert_eq!(entries[1].entry_type, "rate_limit_secondary");
504        assert!(entries[1].limited);
505    }
506
507    #[test]
508    fn codex_usage_entries_adds_summary_when_limit_has_no_windows() {
509        let limit = CodexRateLimit {
510            allowed: false,
511            limit_reached: true,
512            primary_window: None,
513            secondary_window: None,
514        };
515
516        let entries = codex_usage_entries("code_review_rate_limit", &limit);
517
518        assert_eq!(entries.len(), 1);
519        assert_eq!(entries[0].entry_type, "code_review_rate_limit");
520        assert!(entries[0].limited);
521        assert!((entries[0].utilization - 100.0).abs() < f64::EPSILON);
522        assert_eq!(entries[0].resets_at, None);
523    }
524
525    // -----------------------------------------------------------------------
526    // OpenRouter dispatch tests
527    // These tests verify that check_limit() / fetch_status() correctly route
528    // to the openrouter handler when provider == "openrouter", and that a
529    // missing management key causes an immediate error (no HTTP call made).
530    // -----------------------------------------------------------------------
531
532    fn make_openrouter_agent(management_key: Option<&str>) -> Agent {
533        Agent::new(
534            AgentConfig {
535                command: "myai".to_string(),
536                env: None,
537                provider: Some(crate::config::ProviderConfig::Explicit(
538                    "openrouter".to_string(),
539                )),
540                openrouter_management_key: management_key.map(str::to_string),
541                glm_api_key: None,
542            },
543            vec![],
544        )
545    }
546
547    type TestResult = Result<(), Box<dyn std::error::Error>>;
548
549    #[tokio::test(flavor = "current_thread")]
550    async fn check_limit_openrouter_returns_error_when_management_key_is_missing() -> TestResult {
551        // Given: openrouter agent with no management key configured
552        let agent = make_openrouter_agent(None);
553
554        // When: check_limit is called
555        let result = agent.check_limit().await;
556
557        // Then: error mentions the missing key -- no HTTP call should be made
558        let err_msg = result.err().ok_or("expected Err")?.to_string();
559        assert!(err_msg.contains("openrouter_management_key"));
560        Ok(())
561    }
562
563    #[tokio::test(flavor = "current_thread")]
564    async fn fetch_status_openrouter_returns_error_when_management_key_is_missing() -> TestResult {
565        // Given: openrouter agent with no management key configured
566        let agent = make_openrouter_agent(None);
567
568        // When: fetch_status is called
569        let result = agent.fetch_status().await;
570
571        // Then: error mentions the missing key -- no HTTP call should be made
572        let err_msg = result.err().ok_or("expected Err")?.to_string();
573        assert!(err_msg.contains("openrouter_management_key"));
574        Ok(())
575    }
576
577    fn make_api_key_agent(provider: &str) -> Agent {
578        Agent::new(
579            AgentConfig {
580                command: "myai".to_string(),
581                env: None,
582                provider: Some(crate::config::ProviderConfig::Explicit(
583                    provider.to_string(),
584                )),
585                openrouter_management_key: None,
586                glm_api_key: None,
587            },
588            vec![],
589        )
590    }
591
592    // -- zai --
593
594    #[tokio::test(flavor = "current_thread")]
595    async fn check_limit_zai_returns_error_when_api_key_is_missing() -> TestResult {
596        let agent = make_api_key_agent("zai");
597        let result = agent.check_limit().await;
598        let err_msg = result.err().ok_or("expected Err")?.to_string();
599        assert!(
600            err_msg.contains("Z_AI_API_KEY"),
601            "error should mention Z_AI_API_KEY, got: {err_msg}"
602        );
603        Ok(())
604    }
605
606    #[tokio::test(flavor = "current_thread")]
607    async fn fetch_status_zai_returns_error_when_api_key_is_missing() -> TestResult {
608        let agent = make_api_key_agent("zai");
609        let result = agent.fetch_status().await;
610        let err_msg = result.err().ok_or("expected Err")?.to_string();
611        assert!(
612            err_msg.contains("Z_AI_API_KEY"),
613            "error should mention Z_AI_API_KEY, got: {err_msg}"
614        );
615        Ok(())
616    }
617
618    // -- kimi-k2 --
619
620    #[tokio::test(flavor = "current_thread")]
621    async fn check_limit_kimik2_returns_error_when_api_key_is_missing() -> TestResult {
622        let agent = make_api_key_agent("kimi-k2");
623        let result = agent.check_limit().await;
624        let err_msg = result.err().ok_or("expected Err")?.to_string();
625        assert!(
626            err_msg.contains("KIMI_K2_API_KEY"),
627            "error should mention KIMI_K2_API_KEY, got: {err_msg}"
628        );
629        Ok(())
630    }
631
632    #[tokio::test(flavor = "current_thread")]
633    async fn fetch_status_kimik2_returns_error_when_api_key_is_missing() -> TestResult {
634        let agent = make_api_key_agent("kimi-k2");
635        let result = agent.fetch_status().await;
636        let err_msg = result.err().ok_or("expected Err")?.to_string();
637        assert!(
638            err_msg.contains("KIMI_K2_API_KEY"),
639            "error should mention KIMI_K2_API_KEY, got: {err_msg}"
640        );
641        Ok(())
642    }
643
644    // -- warp --
645
646    #[tokio::test(flavor = "current_thread")]
647    async fn check_limit_warp_returns_error_when_api_key_is_missing() -> TestResult {
648        let agent = make_api_key_agent("warp");
649        let result = agent.check_limit().await;
650        let err_msg = result.err().ok_or("expected Err")?.to_string();
651        assert!(
652            err_msg.contains("WARP_API_KEY"),
653            "error should mention WARP_API_KEY, got: {err_msg}"
654        );
655        Ok(())
656    }
657
658    #[tokio::test(flavor = "current_thread")]
659    async fn fetch_status_warp_returns_error_when_api_key_is_missing() -> TestResult {
660        let agent = make_api_key_agent("warp");
661        let result = agent.fetch_status().await;
662        let err_msg = result.err().ok_or("expected Err")?.to_string();
663        assert!(
664            err_msg.contains("WARP_API_KEY"),
665            "error should mention WARP_API_KEY, got: {err_msg}"
666        );
667        Ok(())
668    }
669
670    // -- kiro (CLI-based, no API key needed for dispatch, but must not panic) --
671
672    #[tokio::test(flavor = "current_thread")]
673    async fn check_limit_kiro_returns_error_when_command_not_found() -> TestResult {
674        let agent = make_api_key_agent("kiro");
675        let result = agent.check_limit().await;
676        assert!(result.is_err(), "kiro without CLI should return an error");
677        Ok(())
678    }
679
680    #[tokio::test(flavor = "current_thread")]
681    async fn fetch_status_kiro_returns_error_when_command_not_found() -> TestResult {
682        let agent = make_api_key_agent("kiro");
683        let result = agent.fetch_status().await;
684        assert!(result.is_err(), "kiro without CLI should return an error");
685        Ok(())
686    }
687
688    #[tokio::test(flavor = "current_thread")]
689    async fn check_limit_opencode_go_uses_local_history() -> TestResult {
690        let tmp = tempfile::tempdir()?;
691        let db_path = tmp.path().join("opencode.db");
692        let conn = rusqlite::Connection::open(&db_path)?;
693        conn.execute("CREATE TABLE message (data TEXT NOT NULL)", [])?;
694        conn.execute(
695            "INSERT INTO message (data) VALUES (?1)",
696            [r#"{"role":"assistant","providerID":"opencode-go","cost":6.5,"time":{"completed":4102448400000}}"#],
697        )?;
698        conn.execute(
699            "INSERT INTO message (data) VALUES (?1)",
700            [r#"{"role":"assistant","providerID":"opencode-go","cost":6.0,"time":{"completed":4102461000000}}"#],
701        )?;
702        drop(conn);
703
704        let mut agent = make_api_key_agent("opencode-go");
705        agent.config.env = Some(HashMap::from([(
706            "SEHER_OPENCODE_DB_PATH".to_string(),
707            db_path.display().to_string(),
708        )]));
709        let result = agent.check_limit().await?;
710        assert!(matches!(result, AgentLimit::Limited { .. }));
711        Ok(())
712    }
713
714    #[tokio::test(flavor = "current_thread")]
715    async fn fetch_status_opencode_go_returns_usage_windows() -> TestResult {
716        let tmp = tempfile::tempdir()?;
717        let db_path = tmp.path().join("opencode.db");
718        let conn = rusqlite::Connection::open(&db_path)?;
719        conn.execute("CREATE TABLE message (data TEXT NOT NULL)", [])?;
720        conn.execute(
721            "INSERT INTO message (data) VALUES (?1)",
722            [r#"{"role":"assistant","providerID":"opencode-go","cost":2.25,"time":{"completed":4102461000000}}"#],
723        )?;
724        drop(conn);
725
726        let mut agent = make_api_key_agent("opencode-go");
727        agent.config.env = Some(HashMap::from([(
728            "SEHER_OPENCODE_DB_PATH".to_string(),
729            db_path.display().to_string(),
730        )]));
731        let status = agent.fetch_status().await?;
732        assert_eq!(status.provider.as_deref(), Some("opencode-go"));
733        assert_eq!(status.usage.len(), 3);
734        assert!(
735            status
736                .usage
737                .iter()
738                .any(|entry| entry.entry_type == "five_hour_spend")
739        );
740        assert!(
741            status
742                .usage
743                .iter()
744                .any(|entry| entry.entry_type == "weekly_spend")
745        );
746        assert!(
747            status
748                .usage
749                .iter()
750                .any(|entry| entry.entry_type == "monthly_spend")
751        );
752        Ok(())
753    }
754
755    // -- unknown provider still errors --
756
757    #[tokio::test(flavor = "current_thread")]
758    async fn check_limit_unknown_provider_returns_error() -> TestResult {
759        let agent = make_api_key_agent("nonexistent-provider");
760        let result = agent.check_limit().await;
761        let err_msg = result.err().ok_or("expected Err")?.to_string();
762        assert!(err_msg.contains("Unknown provider"), "got: {err_msg}");
763        Ok(())
764    }
765}
seher/agent/mod.rs

seher/agent/
mod.rs
