edgecrab_types/
usage.rs

1//! Token usage and cost tracking types.
2//!
3//! Normalizes usage across different API response formats
4//! (OpenAI, Anthropic, Codex) into a single structure.
5
6use serde::{Deserialize, Serialize};
7
8use crate::ApiMode;
9
10#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
11pub struct Usage {
12    pub input_tokens: u64,
13    pub output_tokens: u64,
14    pub cache_read_tokens: u64,
15    pub cache_write_tokens: u64,
16    pub reasoning_tokens: u64,
17    pub total_tokens: u64,
18}
19
20#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
21pub struct Cost {
22    pub input_cost: f64,
23    pub output_cost: f64,
24    pub cache_read_cost: f64,
25    pub cache_write_cost: f64,
26    pub total_cost: f64,
27}
28
29impl Usage {
30    /// Compute total tokens from components.
31    pub fn compute_total(&mut self) {
32        self.total_tokens = self.input_tokens + self.output_tokens;
33    }
34}
35
36/// Normalize raw API usage JSON into our unified Usage struct.
37///
38/// Each provider returns usage in a different shape — this function
39/// maps them all into one representation.
40pub fn normalize_usage(raw: &serde_json::Value, api_mode: ApiMode) -> Usage {
41    /// Extract a u64 token count from a JSON field, defaulting to 0.
42    fn tok(v: &serde_json::Value) -> u64 {
43        v.as_u64().unwrap_or(0)
44    }
45
46    match api_mode {
47        ApiMode::ChatCompletions => Usage {
48            input_tokens: tok(&raw["prompt_tokens"]),
49            output_tokens: tok(&raw["completion_tokens"]),
50            cache_read_tokens: tok(&raw["prompt_tokens_details"]["cached_tokens"]),
51            total_tokens: tok(&raw["total_tokens"]),
52            ..Default::default()
53        },
54        ApiMode::AnthropicMessages => {
55            let input = tok(&raw["input_tokens"]);
56            let output = tok(&raw["output_tokens"]);
57            Usage {
58                input_tokens: input,
59                output_tokens: output,
60                cache_read_tokens: tok(&raw["cache_read_input_tokens"]),
61                cache_write_tokens: tok(&raw["cache_creation_input_tokens"]),
62                total_tokens: input + output,
63                ..Default::default()
64            }
65        }
66        ApiMode::CodexResponses => {
67            let input = tok(&raw["input_tokens"]);
68            let output = tok(&raw["output_tokens"]);
69            Usage {
70                input_tokens: input,
71                output_tokens: output,
72                total_tokens: input + output,
73                ..Default::default()
74            }
75        }
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn normalize_chat_completions() {
85        let raw = serde_json::json!({
86            "prompt_tokens": 100,
87            "completion_tokens": 50,
88            "total_tokens": 150,
89            "prompt_tokens_details": { "cached_tokens": 20 }
90        });
91        let usage = normalize_usage(&raw, ApiMode::ChatCompletions);
92        assert_eq!(usage.input_tokens, 100);
93        assert_eq!(usage.output_tokens, 50);
94        assert_eq!(usage.cache_read_tokens, 20);
95        assert_eq!(usage.total_tokens, 150);
96    }
97
98    #[test]
99    fn normalize_anthropic() {
100        let raw = serde_json::json!({
101            "input_tokens": 200,
102            "output_tokens": 80,
103            "cache_read_input_tokens": 50,
104            "cache_creation_input_tokens": 10
105        });
106        let usage = normalize_usage(&raw, ApiMode::AnthropicMessages);
107        assert_eq!(usage.input_tokens, 200);
108        assert_eq!(usage.output_tokens, 80);
109        assert_eq!(usage.cache_read_tokens, 50);
110        assert_eq!(usage.cache_write_tokens, 10);
111        assert_eq!(usage.total_tokens, 280);
112    }
113
114    #[test]
115    fn normalize_codex() {
116        let raw = serde_json::json!({
117            "input_tokens": 300,
118            "output_tokens": 100
119        });
120        let usage = normalize_usage(&raw, ApiMode::CodexResponses);
121        assert_eq!(usage.input_tokens, 300);
122        assert_eq!(usage.output_tokens, 100);
123        assert_eq!(usage.total_tokens, 400);
124    }
125
126    #[test]
127    fn usage_roundtrip() {
128        let usage = Usage {
129            input_tokens: 100,
130            output_tokens: 50,
131            cache_read_tokens: 20,
132            cache_write_tokens: 5,
133            reasoning_tokens: 10,
134            total_tokens: 150,
135        };
136        let json = serde_json::to_string(&usage).expect("serialize");
137        let deser: Usage = serde_json::from_str(&json).expect("deserialize");
138        assert_eq!(usage, deser);
139    }
140
141    #[test]
142    fn cost_roundtrip() {
143        let cost = Cost {
144            input_cost: 0.001,
145            output_cost: 0.003,
146            total_cost: 0.004,
147            ..Default::default()
148        };
149        let json = serde_json::to_string(&cost).expect("serialize");
150        let deser: Cost = serde_json::from_str(&json).expect("deserialize");
151        assert_eq!(cost, deser);
152    }
153}
edgecrab_types/usage.rs

edgecrab_types/
usage.rs
