aacode_rs/
usage.rs

1// Copyright (c) 2026 xiefujin <490021684@qq.com>
2// Licensed under GPL-3.0, see LICENSE file for full license terms.
3
4//! Append-only provider-usage recording.
5//!
6//! Writes one JSONL line per model call under `.aacode/usage/usage.jsonl`.
7//! Only token counts and request metadata are stored — never prompt or
8//! response text. Telemetry must never break a task, so all errors are
9//! swallowed.
10//!
11//! Env: `AACODE_USAGE_LOG=0` disables; `AACODE_USAGE_VERBOSE=1` prints one
12//! compact line per call.
13
14use serde_json::Value;
15use std::io::Write;
16use std::path::{Path, PathBuf};
17
18pub struct UsageTracker {
19    project_path: PathBuf,
20    enabled: bool,
21    pub verbose: bool,
22}
23
24impl UsageTracker {
25    pub fn new(project_path: &Path) -> Self {
26        let enabled = std::env::var("AACODE_USAGE_LOG")
27            .map(|v| v != "0")
28            .unwrap_or(true);
29        let verbose = std::env::var("AACODE_USAGE_VERBOSE")
30            .map(|v| v == "1")
31            .unwrap_or(false);
32        UsageTracker {
33            project_path: project_path.to_path_buf(),
34            enabled,
35            verbose,
36        }
37    }
38
39    /// Append one event (a `ts` field is added). Returns false when disabled or
40    /// on any I/O error.
41    pub fn record(&self, event: &Value) -> bool {
42        if !self.enabled {
43            return false;
44        }
45        let dir = self.project_path.join(".aacode").join("usage");
46        if std::fs::create_dir_all(&dir).is_err() {
47            return false;
48        }
49        let path = dir.join("usage.jsonl");
50        let mut line = event.clone();
51        if let Some(obj) = line.as_object_mut() {
52            obj.insert("ts".into(), Value::String(crate::session::now_iso_ms()));
53        }
54        let Ok(mut f) = std::fs::OpenOptions::new()
55            .create(true)
56            .append(true)
57            .open(&path)
58        else {
59            return false;
60        };
61        f.write_all(line.to_string().as_bytes()).is_ok() && f.write_all(b"\n").is_ok()
62    }
63}
64
65/// Build the JSONL event for one model call. `usage` is the normalized usage
66/// (if the provider reported it); pass `None` when missing.
67#[allow(clippy::too_many_arguments)]
68pub fn usage_event(
69    session_id: &str,
70    run_id: &str,
71    iteration: u32,
72    gateway: &str,
73    model: &str,
74    estimated_input_tokens: usize,
75    messages_count: usize,
76    compacted: bool,
77    usage: Option<&crate::llm::Usage>,
78) -> Value {
79    let mut event = serde_json::json!({
80        "source": "main",
81        "session_id": session_id,
82        "run_id": run_id,
83        "iteration": iteration,
84        "gateway": gateway,
85        "model": model,
86        "stream": true,
87        "status": "ok",
88        "usage_source": if usage.is_some() { "provider" } else { "missing" },
89        "estimated_input_tokens": estimated_input_tokens,
90        "context_tokens_estimate": estimated_input_tokens,
91        "messages_count": messages_count,
92        "compacted": compacted,
93    });
94    if let Some(u) = usage {
95        event["input_tokens"] = serde_json::json!(u.input_tokens);
96        event["output_tokens"] = serde_json::json!(u.output_tokens);
97        event["total_tokens"] = serde_json::json!(u.total_tokens);
98        event["cache_read_tokens"] = serde_json::json!(u.cache_read_tokens);
99        event["cache_write_tokens"] = serde_json::json!(u.cache_write_tokens);
100        event["cache_miss_tokens"] = serde_json::json!(u.cache_miss_tokens);
101        event["reasoning_tokens"] = serde_json::json!(u.reasoning_tokens);
102        event["cache_hit_rate"] = serde_json::json!(u.cache_hit_rate());
103    }
104    event
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use crate::llm::Usage;
111
112    #[test]
113    fn event_includes_cache_hit_rate() {
114        let u = Usage {
115            input_tokens: 1000,
116            output_tokens: 100,
117            total_tokens: 1100,
118            cache_read_tokens: 900,
119            cache_write_tokens: 0,
120            cache_miss_tokens: 100,
121            reasoning_tokens: 5,
122        };
123        let ev = usage_event("s", "r", 1, "openai", "m", 999, 10, false, Some(&u));
124        assert_eq!(ev["usage_source"], "provider");
125        assert!((ev["cache_hit_rate"].as_f64().unwrap() - 0.9).abs() < 1e-9);
126        assert_eq!(ev["input_tokens"], 1000);
127    }
128
129    #[test]
130    fn event_without_usage_is_marked_missing() {
131        let ev = usage_event("s", "r", 1, "openai", "m", 10, 2, false, None);
132        assert_eq!(ev["usage_source"], "missing");
133        assert!(ev.get("cache_hit_rate").is_none());
134    }
135
136    #[test]
137    fn records_jsonl_line() {
138        let dir = std::env::temp_dir().join(format!(
139            "aacode_usage_{}_{}",
140            std::process::id(),
141            uuid::Uuid::new_v4().simple()
142        ));
143        std::fs::create_dir_all(&dir).unwrap();
144        let tracker = UsageTracker::new(&dir);
145        let ev = usage_event("s", "r", 1, "openai", "m", 10, 2, false, None);
146        assert!(tracker.record(&ev));
147        let text =
148            std::fs::read_to_string(dir.join(".aacode").join("usage").join("usage.jsonl")).unwrap();
149        assert!(text.contains("\"usage_source\":\"missing\""), "{text}");
150        assert!(text.contains("\"ts\""), "{text}");
151        assert_eq!(text.lines().count(), 1);
152    }
153}
aacode_rs/usage.rs

aacode_rs/
usage.rs
