use std::time::Duration;
use anyhow::{Context, Result};
use tokio::sync::watch;
use tracing::{debug, info, warn};
use super::parser::parse_usage_output;
use super::types::UsageSnapshot;
use crate::tmux::TmuxClient;
pub type UsageSnapshotSender = watch::Sender<UsageSnapshot>;
pub type UsageSnapshotReceiver = watch::Receiver<UsageSnapshot>;
pub fn usage_channel() -> (UsageSnapshotSender, UsageSnapshotReceiver) {
watch::channel(UsageSnapshot::default())
}
pub async fn fetch_usage(session: &str) -> Result<UsageSnapshot> {
let tmux = TmuxClient::new();
let home = std::env::var("HOME")
.context("HOME environment variable is not set; cannot determine trusted directory")?;
let target = tmux
.new_window(session, &home, Some("tmai-usage"))
.context("Failed to create hidden window for usage fetch")?;
info!("Usage fetch: created hidden pane {}", target);
tmux.run_command(&target, "claude")
.context("Failed to start claude in usage pane")?;
let started = wait_for_claude_ready(&tmux, &target, Duration::from_secs(30)).await;
if !started {
let _ = tmux.kill_pane(&target);
anyhow::bail!("Claude Code did not start within timeout");
}
debug!("Usage fetch: Claude Code ready, sending /usage");
tokio::time::sleep(Duration::from_millis(500)).await;
tmux.send_keys_literal(&target, "/usage")
.context("Failed to send /usage")?;
tokio::time::sleep(Duration::from_millis(300)).await;
tmux.send_keys(&target, "Enter")
.context("Failed to send Enter after /usage")?;
let usage_text = wait_for_usage_output(&tmux, &target, Duration::from_secs(15)).await;
let _ = tmux.send_keys(&target, "Escape");
tokio::time::sleep(Duration::from_millis(300)).await;
let _ = tmux.send_keys(&target, "C-c");
tokio::time::sleep(Duration::from_millis(300)).await;
let _ = tmux.send_keys(&target, "C-c");
tokio::time::sleep(Duration::from_millis(500)).await;
let _ = tmux.kill_pane(&target);
info!("Usage fetch: cleaned up pane {}", target);
match usage_text {
Some(text) => {
let snapshot = parse_usage_output(&text);
if snapshot.meters.is_empty() {
anyhow::bail!("Failed to parse usage output (no meters found)");
}
Ok(snapshot)
}
None => anyhow::bail!("Usage overlay did not appear within timeout"),
}
}
async fn wait_for_claude_ready(tmux: &TmuxClient, target: &str, timeout: Duration) -> bool {
let start = std::time::Instant::now();
let poll_interval = Duration::from_millis(500);
let mut trust_confirmed = false;
while start.elapsed() < timeout {
tokio::time::sleep(poll_interval).await;
if let Ok(content) = tmux.capture_pane_plain(target) {
if content.contains("-- INSERT --") {
debug!("Usage fetch: detected INSERT mode, Claude is ready");
return true;
}
if !trust_confirmed && content.contains("Yes, I trust this folder") {
debug!("Usage fetch: auto-confirming trust prompt");
let _ = tmux.send_keys(target, "Enter");
trust_confirmed = true;
}
}
}
warn!("Usage fetch: timed out waiting for Claude Code to start");
false
}
async fn wait_for_usage_output(
tmux: &TmuxClient,
target: &str,
timeout: Duration,
) -> Option<String> {
let start = std::time::Instant::now();
let poll_interval = Duration::from_millis(300);
while start.elapsed() < timeout {
tokio::time::sleep(poll_interval).await;
if let Ok(content) = tmux.capture_pane_plain(target) {
if content.contains("% used") {
debug!("Usage fetch: detected usage output");
return Some(content);
}
debug!(
"Usage fetch: waiting for /usage output ({:.1}s elapsed)",
start.elapsed().as_secs_f32()
);
}
}
warn!("Usage fetch: timed out waiting for /usage output");
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_usage_channel() {
let (tx, rx) = usage_channel();
assert!(rx.borrow().meters.is_empty());
let snapshot = UsageSnapshot {
meters: vec![],
fetched_at: Some(chrono::Utc::now()),
fetching: false,
error: None,
};
tx.send(snapshot).unwrap();
assert!(rx.borrow().fetched_at.is_some());
}
}