inference_gateway_adk/server/
usage_tracker.rs

1use inference_gateway_sdk::CompletionUsage;
2use serde_json::{Map, Value, json};
3use std::sync::Mutex;
4
5/// Accumulates token usage and execution statistics over a single task run.
6///
7/// The default task handlers create one [`UsageTracker`] per task, thread it
8/// through [`run_tool_loop`](super::task_handler), and - when usage metadata is
9/// enabled - merge [`metadata`](Self::metadata) into the task's
10/// `metadata` field once the task reaches a terminal state.
11///
12/// Counters are split into two groups:
13/// - **Token usage** (`prompt_tokens` / `completion_tokens` / `total_tokens`),
14///   summed from each [`CompletionUsage`] the gateway returns. Every
15///   [`add_token_usage`](Self::add_token_usage) call also bumps an internal
16///   `llm_calls` counter so [`metadata`](Self::metadata) can decide whether a
17///   `usage` block is meaningful.
18/// - **Execution stats** (`iterations` / `messages` / `tool_calls` /
19///   `failed_tools`), counted as the agent loop runs.
20///
21/// All mutators take `&self` (interior mutability via a [`Mutex`]) so the
22/// tracker can be shared behind a shared reference across `.await` points.
23#[derive(Debug, Default)]
24pub struct UsageTracker {
25    inner: Mutex<Stats>,
26}
27
28#[derive(Debug, Default, Clone)]
29struct Stats {
30    prompt_tokens: i64,
31    completion_tokens: i64,
32    total_tokens: i64,
33    iterations: u64,
34    messages: u64,
35    tool_calls: u64,
36    failed_tools: u64,
37    llm_calls: u64,
38}
39
40impl UsageTracker {
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    fn lock(&self) -> std::sync::MutexGuard<'_, Stats> {
46        self.inner.lock().unwrap_or_else(|e| e.into_inner())
47    }
48
49    /// Add the token counts from one chat-completion response and record that
50    /// an LLM call was made.
51    pub fn add_token_usage(&self, usage: &CompletionUsage) {
52        let mut s = self.lock();
53        s.prompt_tokens += usage.prompt_tokens;
54        s.completion_tokens += usage.completion_tokens;
55        s.total_tokens += usage.total_tokens;
56        s.llm_calls += 1;
57    }
58
59    /// Count one pass through the agent's model ↔ tool loop.
60    pub fn increment_iteration(&self) {
61        self.lock().iterations += 1;
62    }
63
64    /// Count `count` conversation messages produced during the loop (the
65    /// default handlers use this for tool-result messages fed back to the
66    /// model).
67    pub fn add_messages(&self, count: usize) {
68        self.lock().messages += count as u64;
69    }
70
71    /// Count one dispatched tool call.
72    pub fn increment_tool_calls(&self) {
73        self.lock().tool_calls += 1;
74    }
75
76    /// Count one tool call that failed (missing handler or handler error).
77    pub fn increment_failed_tools(&self) {
78        self.lock().failed_tools += 1;
79    }
80
81    /// Whether anything worth reporting was recorded. Mirrors the Go ADK:
82    /// true when at least one LLM call, iteration, message, or tool call
83    /// happened.
84    pub fn has_usage(&self) -> bool {
85        let s = self.lock();
86        s.llm_calls > 0 || s.iterations > 0 || s.messages > 0 || s.tool_calls > 0
87    }
88
89    /// Build the metadata object merged into a terminal task. Always emits an
90    /// `execution_stats` block; only emits a `usage` block when at least one
91    /// LLM call contributed token counts (otherwise the zeros would be
92    /// misleading).
93    pub fn metadata(&self) -> Map<String, Value> {
94        let s = self.lock();
95        let mut map = Map::new();
96        if s.llm_calls > 0 {
97            map.insert(
98                "usage".to_string(),
99                json!({
100                    "prompt_tokens": s.prompt_tokens,
101                    "completion_tokens": s.completion_tokens,
102                    "total_tokens": s.total_tokens,
103                }),
104            );
105        }
106        map.insert(
107            "execution_stats".to_string(),
108            json!({
109                "iterations": s.iterations,
110                "messages": s.messages,
111                "tool_calls": s.tool_calls,
112                "failed_tools": s.failed_tools,
113            }),
114        );
115        map
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    fn usage(prompt: i64, completion: i64, total: i64) -> CompletionUsage {
124        CompletionUsage {
125            prompt_tokens: prompt,
126            completion_tokens: completion,
127            total_tokens: total,
128        }
129    }
130
131    #[test]
132    fn fresh_tracker_reports_no_usage() {
133        let tracker = UsageTracker::new();
134        assert!(!tracker.has_usage());
135        // Only execution_stats is present, with all-zero counters.
136        let meta = tracker.metadata();
137        assert!(!meta.contains_key("usage"));
138        let stats = &meta["execution_stats"];
139        assert_eq!(stats["iterations"], 0);
140        assert_eq!(stats["messages"], 0);
141        assert_eq!(stats["tool_calls"], 0);
142        assert_eq!(stats["failed_tools"], 0);
143    }
144
145    #[test]
146    fn token_usage_accumulates_across_calls() {
147        let tracker = UsageTracker::new();
148        tracker.add_token_usage(&usage(10, 5, 15));
149        tracker.add_token_usage(&usage(20, 8, 28));
150
151        assert!(tracker.has_usage());
152        let meta = tracker.metadata();
153        let usage_block = &meta["usage"];
154        assert_eq!(usage_block["prompt_tokens"], 30);
155        assert_eq!(usage_block["completion_tokens"], 13);
156        assert_eq!(usage_block["total_tokens"], 43);
157    }
158
159    #[test]
160    fn execution_stats_count_loop_activity() {
161        let tracker = UsageTracker::new();
162        tracker.increment_iteration();
163        tracker.increment_iteration();
164        tracker.add_messages(3);
165        tracker.increment_tool_calls();
166        tracker.increment_tool_calls();
167        tracker.increment_tool_calls();
168        tracker.increment_failed_tools();
169
170        assert!(tracker.has_usage());
171        let stats = &tracker.metadata()["execution_stats"];
172        assert_eq!(stats["iterations"], 2);
173        assert_eq!(stats["messages"], 3);
174        assert_eq!(stats["tool_calls"], 3);
175        assert_eq!(stats["failed_tools"], 1);
176    }
177
178    #[test]
179    fn iterations_alone_count_as_usage_without_a_usage_block() {
180        // A streaming turn that produced no token counts still reports
181        // execution stats but omits the usage block.
182        let tracker = UsageTracker::new();
183        tracker.increment_iteration();
184
185        assert!(tracker.has_usage());
186        let meta = tracker.metadata();
187        assert!(!meta.contains_key("usage"));
188        assert_eq!(meta["execution_stats"]["iterations"], 1);
189    }
190
191    #[test]
192    fn metadata_matches_expected_shape() {
193        let tracker = UsageTracker::new();
194        tracker.add_token_usage(&usage(123, 45, 168));
195        tracker.increment_iteration();
196        tracker.increment_iteration();
197        tracker.add_messages(7);
198        tracker.increment_tool_calls();
199        tracker.increment_tool_calls();
200        tracker.increment_tool_calls();
201        tracker.increment_failed_tools();
202
203        let meta = Value::Object(tracker.metadata());
204        let expected = json!({
205            "usage": {
206                "prompt_tokens": 123,
207                "completion_tokens": 45,
208                "total_tokens": 168,
209            },
210            "execution_stats": {
211                "iterations": 2,
212                "messages": 7,
213                "tool_calls": 3,
214                "failed_tools": 1,
215            },
216        });
217        assert_eq!(meta, expected);
218    }
219}
inference_gateway_adk/server/usage_tracker.rs

inference_gateway_adk/server/
usage_tracker.rs
