bamboo_engine/
token_usage_log.rs

1//! Append-only, per-LLM-call token-usage log.
2//!
3//! Each agent LLM call appends one [`TokenUsageRecord`] (as a JSON line) to the
4//! session's `token-usage.jsonl`, written into the session's own directory next
5//! to `session.json`. Unlike `session.token_usage` — a single snapshot that is
6//! overwritten on every call — this log keeps the full per-round history so
7//! cache effectiveness and cost can be analyzed offline: confirming the 1-hour
8//! prompt-cache TTL keeps hitting across pauses, or measuring the cold round
9//! right after a context compression.
10//!
11//! The record bridges two sources: the prompt-side budget snapshot
12//! ([`TokenBudgetUsage`], already on the session) and the server-returned stream
13//! stats — notably cache-creation and OpenAI cache-write dimensions, which are
14//! NOT part of `TokenBudgetUsage` and would otherwise only exist transiently.
15
16use bamboo_domain::TokenBudgetUsage;
17use serde::Serialize;
18
19/// One per-LLM-call usage record. Flattened for line-oriented analysis
20/// (jq / DuckDB / pandas over the JSONL).
21#[derive(Debug, Clone, Serialize)]
22pub struct TokenUsageRecord {
23    /// RFC3339 wall-clock timestamp captured when the call completed.
24    pub ts: String,
25    pub session_id: String,
26    pub model: String,
27    pub provider: String,
28    /// Conversation length (message count) at emit time — a monotonic-ish
29    /// ordinal to correlate records with conversation growth.
30    pub message_count: usize,
31
32    // --- server-returned usage (this call) ---
33    pub cache_creation_input_tokens: u64,
34    pub cache_read_input_tokens: u64,
35    /// OpenAI cache-write volume. This can overlap the fresh-input count and
36    /// therefore must not be added to the disjoint prompt-total equation.
37    pub cache_write_input_tokens: u64,
38    /// Non-cached "fresh" input tokens (server-reported), disjoint from the two
39    /// cache counts. The precise prompt size is
40    /// `input_tokens + cache_read + cache_creation`, and the exact cache-hit
41    /// ratio is `cache_read / that_sum`.
42    pub input_tokens: u64,
43    /// Total provider output. For reasoning-capable OpenAI-compatible models,
44    /// `thinking_tokens` is a subset breakdown of this value, not an additional
45    /// quantity.
46    pub output_tokens: u64,
47    pub thinking_tokens: u64,
48
49    // --- prompt-side budget snapshot (this call) ---
50    pub system_tokens: u32,
51    pub summary_tokens: u32,
52    pub window_tokens: u32,
53    pub total_tokens: u32,
54    pub max_context_tokens: u32,
55    pub budget_limit: u32,
56    pub prompt_cached_tool_outputs: usize,
57    pub prompt_cached_tool_tokens_saved: u32,
58    pub truncation_occurred: bool,
59    pub segments_removed: usize,
60}
61
62impl TokenUsageRecord {
63    /// Build a record from the prompt-side budget snapshot (`usage`) and the
64    /// server-side stream stats. Cache creation and cache writes live only on
65    /// the stream output — they are not part of [`TokenBudgetUsage`] — so they
66    /// are passed in explicitly.
67    #[allow(clippy::too_many_arguments)]
68    pub fn new(
69        ts: String,
70        session_id: &str,
71        model: &str,
72        provider: &str,
73        message_count: usize,
74        usage: Option<&TokenBudgetUsage>,
75        cache_creation_input_tokens: u64,
76        cache_read_input_tokens: u64,
77        cache_write_input_tokens: u64,
78        input_tokens: u64,
79        output_tokens: u64,
80        thinking_tokens: u64,
81    ) -> Self {
82        Self {
83            ts,
84            session_id: session_id.to_string(),
85            model: model.to_string(),
86            provider: provider.to_string(),
87            message_count,
88            cache_creation_input_tokens,
89            cache_read_input_tokens,
90            cache_write_input_tokens,
91            input_tokens,
92            output_tokens,
93            thinking_tokens,
94            system_tokens: usage.map(|u| u.system_tokens).unwrap_or(0),
95            summary_tokens: usage.map(|u| u.summary_tokens).unwrap_or(0),
96            window_tokens: usage.map(|u| u.window_tokens).unwrap_or(0),
97            total_tokens: usage.map(|u| u.total_tokens).unwrap_or(0),
98            max_context_tokens: usage.map(|u| u.max_context_tokens).unwrap_or(0),
99            budget_limit: usage.map(|u| u.budget_limit).unwrap_or(0),
100            prompt_cached_tool_outputs: usage.map(|u| u.prompt_cached_tool_outputs).unwrap_or(0),
101            prompt_cached_tool_tokens_saved: usage
102                .map(|u| u.prompt_cached_tool_tokens_saved)
103                .unwrap_or(0),
104            truncation_occurred: usage.map(|u| u.truncation_occurred).unwrap_or(false),
105            segments_removed: usage.map(|u| u.segments_removed).unwrap_or(0),
106        }
107    }
108
109    /// Serialize to a single-line JSON string (no trailing newline; the storage
110    /// layer frames the line).
111    pub fn to_json_line(&self) -> Result<String, serde_json::Error> {
112        serde_json::to_string(self)
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn record_serializes_to_single_json_line_with_cache_creation() {
122        let usage = TokenBudgetUsage {
123            system_tokens: 5000,
124            summary_tokens: 2000,
125            window_tokens: 3000,
126            total_tokens: 10000,
127            max_context_tokens: 200_000,
128            budget_limit: 180_000,
129            truncation_occurred: false,
130            segments_removed: 0,
131            prompt_cached_tool_outputs: 1,
132            prompt_cached_tool_tokens_saved: 42,
133            thinking_tokens: 7,
134            cache_read_input_tokens: 12_000,
135        };
136        let record = TokenUsageRecord::new(
137            "2026-06-15T00:00:00Z".to_string(),
138            "sess-1",
139            "claude-opus-4-8",
140            "anthropic",
141            24,
142            Some(&usage),
143            1500, // cache_creation — only present on the stream output
144            12_000,
145            0,
146            800, // input_tokens (fresh, non-cached)
147            300,
148            7,
149        );
150        let line = record.to_json_line().expect("serializes");
151        assert!(!line.contains('\n'), "must be a single line");
152        assert!(line.contains("\"cache_creation_input_tokens\":1500"));
153        assert!(line.contains("\"cache_read_input_tokens\":12000"));
154        assert!(line.contains("\"cache_write_input_tokens\":0"));
155        assert!(line.contains("\"input_tokens\":800"));
156        assert!(line.contains("\"session_id\":\"sess-1\""));
157    }
158
159    #[test]
160    fn openai_flat_log_contract_keeps_fresh_cache_and_total_disjoint() {
161        let record = TokenUsageRecord::new(
162            "2026-07-29T00:00:00Z".to_string(),
163            "sess-openai",
164            "gpt-5",
165            "openai",
166            4,
167            None,
168            0,
169            768,
170            64,
171            232,
172            120,
173            20,
174        );
175
176        let prompt_total = record
177            .input_tokens
178            .saturating_add(record.cache_read_input_tokens)
179            .saturating_add(record.cache_creation_input_tokens);
180        let cache_hit_ratio = record.cache_read_input_tokens as f64 / prompt_total as f64;
181
182        assert_eq!(record.input_tokens, 232);
183        assert_eq!(prompt_total, 1000);
184        assert_eq!(record.cache_write_input_tokens, 64);
185        assert!((cache_hit_ratio - 0.768).abs() < f64::EPSILON);
186        assert_eq!(record.output_tokens, 120);
187        assert_eq!(record.thinking_tokens, 20);
188    }
189}
bamboo_engine/token_usage_log.rs

bamboo_engine/
token_usage_log.rs
