apollo/tools/
claude_usage.rs

1//! Claude usage tracking tool — check rate limits and remaining quota
2
3use async_trait::async_trait;
4use serde_json::json;
5use std::sync::Arc;
6
7use super::traits::{Tool, ToolResult, ToolSpec};
8use crate::cost::CostTracker;
9
10pub struct ClaudeUsageTool {
11    cost_tracker: Arc<CostTracker>,
12}
13
14impl ClaudeUsageTool {
15    pub fn new(cost_tracker: Arc<CostTracker>) -> Self {
16        Self { cost_tracker }
17    }
18}
19
20#[async_trait]
21impl Tool for ClaudeUsageTool {
22    fn name(&self) -> &str {
23        "claude_usage"
24    }
25
26    fn spec(&self) -> ToolSpec {
27        ToolSpec {
28            name: "claude_usage".to_string(),
29            description: "Check Claude API usage, rate limits, and remaining quota. Shows requests/tokens remaining and cost summary.".to_string(),
30            parameters: json!({
31                "type": "object",
32                "properties": {
33                    "action": {
34                        "type": "string",
35                        "enum": ["limits", "cost", "both"],
36                        "description": "What to check: 'limits' (rate limits), 'cost' (spending), or 'both'"
37                    }
38                },
39                "required": ["action"]
40            }),
41        }
42    }
43
44    async fn execute(&self, arguments: &str) -> anyhow::Result<ToolResult> {
45        let args: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
46        let action = args["action"].as_str().unwrap_or("both");
47
48        let mut output = Vec::new();
49
50        // Rate limits
51        if action == "limits" || action == "both" {
52            if let Some(limits) = self.cost_tracker.get_rate_limits().await {
53                output.push("⚡ Claude API Rate Limits\n".to_string());
54
55                if let (Some(rem), Some(lim)) = (limits.requests_remaining, limits.requests_limit) {
56                    let pct = (rem as f64 / lim as f64 * 100.0) as usize;
57                    output.push(format!("Requests: {}/{} ({}% left)", rem, lim, pct));
58                }
59
60                if let (Some(rem), Some(lim)) =
61                    (limits.input_tokens_remaining, limits.input_tokens_limit)
62                {
63                    let pct = (rem as f64 / lim as f64 * 100.0) as usize;
64                    output.push(format!(
65                        "Input tokens: {}/{} ({}% left)",
66                        format_tokens(rem),
67                        format_tokens(lim),
68                        pct
69                    ));
70                }
71
72                if let (Some(rem), Some(lim)) =
73                    (limits.output_tokens_remaining, limits.output_tokens_limit)
74                {
75                    let pct = (rem as f64 / lim as f64 * 100.0) as usize;
76                    output.push(format!(
77                        "Output tokens: {}/{} ({}% left)",
78                        format_tokens(rem),
79                        format_tokens(lim),
80                        pct
81                    ));
82                }
83
84                if let Some(reset) = limits.tokens_reset {
85                    output.push(format!("Resets: {}", reset));
86                }
87
88                output.push(String::new());
89            } else {
90                output.push(
91                    "⚠️ No rate limit data yet (need at least one API call first)".to_string(),
92                );
93                output.push(String::new());
94            }
95        }
96
97        // Cost summary
98        if action == "cost" || action == "both" {
99            let summary = self.cost_tracker.summary().await;
100
101            output.push("💰 Cost Summary\n".to_string());
102            output.push(format!("Total spent: ${:.4}", summary.total_cost));
103            output.push(format!(
104                "Total tokens: {}",
105                format_tokens(summary.total_tokens)
106            ));
107            output.push(format!("API calls: {}", summary.call_count));
108
109            if !summary.by_model.is_empty() {
110                output.push(String::new());
111                output.push("By model:".to_string());
112                let mut models: Vec<_> = summary.by_model.iter().collect();
113                models.sort_by(|a, b| b.1.partial_cmp(a.1).unwrap());
114
115                for (model, cost) in models {
116                    output.push(format!("• {}: ${:.4}", model, cost));
117                }
118            }
119        }
120
121        Ok(ToolResult::success(output.join("\n")))
122    }
123}
124
125fn format_tokens(n: usize) -> String {
126    if n >= 1_000_000 {
127        format!("{:.1}M", n as f64 / 1_000_000.0)
128    } else if n >= 1_000 {
129        format!("{:.1}K", n as f64 / 1_000.0)
130    } else {
131        n.to_string()
132    }
133}
apollo/tools/claude_usage.rs

apollo/tools/
claude_usage.rs
