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/// Non-content context-management correlation captured for one provider call.
20///
21/// These fields deliberately contain only strategy/lifecycle identifiers. Raw
22/// messages, search queries/results, memory, tool arguments, file paths, and
23/// provider payloads must never enter the token-usage sidecar.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct ContextManagementTelemetry {
26    pub strategy: String,
27    pub model_context_epoch: u64,
28    pub model_context_reset_reason: Option<String>,
29    pub retrieval_archive_event_count: usize,
30    pub latest_retrieval_archive_event_id: Option<String>,
31    pub latest_retrieval_archive_trigger_type: Option<String>,
32}
33
34impl Default for ContextManagementTelemetry {
35    fn default() -> Self {
36        Self {
37            strategy: "summary".to_string(),
38            model_context_epoch: 0,
39            model_context_reset_reason: None,
40            retrieval_archive_event_count: 0,
41            latest_retrieval_archive_event_id: None,
42            latest_retrieval_archive_trigger_type: None,
43        }
44    }
45}
46
47/// One per-LLM-call usage record. Flattened for line-oriented analysis
48/// (jq / DuckDB / pandas over the JSONL).
49#[derive(Debug, Clone, Serialize)]
50pub struct TokenUsageRecord {
51    /// RFC3339 wall-clock timestamp captured when the call completed.
52    pub ts: String,
53    pub session_id: String,
54    pub model: String,
55    pub provider: String,
56    /// Conversation length (message count) at emit time — a monotonic-ish
57    /// ordinal to correlate records with conversation growth.
58    pub message_count: usize,
59
60    // --- server-returned usage (this call) ---
61    pub cache_creation_input_tokens: u64,
62    pub cache_read_input_tokens: u64,
63    /// OpenAI cache-write volume. This can overlap the fresh-input count and
64    /// therefore must not be added to the disjoint prompt-total equation.
65    pub cache_write_input_tokens: u64,
66    /// Non-cached "fresh" input tokens (server-reported), disjoint from the two
67    /// cache counts. The precise prompt size is
68    /// `input_tokens + cache_read + cache_creation`, and the exact cache-hit
69    /// ratio is `cache_read / that_sum`.
70    pub input_tokens: u64,
71    /// Total provider output. For reasoning-capable OpenAI-compatible models,
72    /// `thinking_tokens` is a subset breakdown of this value, not an additional
73    /// quantity.
74    pub output_tokens: u64,
75    pub thinking_tokens: u64,
76
77    // --- prompt-side budget snapshot (this call) ---
78    pub system_tokens: u32,
79    pub summary_tokens: u32,
80    pub window_tokens: u32,
81    pub total_tokens: u32,
82    pub max_context_tokens: u32,
83    pub budget_limit: u32,
84    pub prompt_cached_tool_outputs: usize,
85    pub prompt_cached_tool_tokens_saved: u32,
86    pub truncation_occurred: bool,
87    pub segments_removed: usize,
88
89    // --- context-management correlation (never raw content) ---
90    pub context_management_strategy: String,
91    pub model_context_epoch: u64,
92    #[serde(skip_serializing_if = "Option::is_none")]
93    pub model_context_reset_reason: Option<String>,
94    pub retrieval_archive_event_count: usize,
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub latest_retrieval_archive_event_id: Option<String>,
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub latest_retrieval_archive_trigger_type: Option<String>,
99}
100
101impl TokenUsageRecord {
102    /// Build a record from the prompt-side budget snapshot (`usage`) and the
103    /// server-side stream stats. Cache creation and cache writes live only on
104    /// the stream output — they are not part of [`TokenBudgetUsage`] — so they
105    /// are passed in explicitly.
106    #[allow(clippy::too_many_arguments)]
107    pub fn new(
108        ts: String,
109        session_id: &str,
110        model: &str,
111        provider: &str,
112        message_count: usize,
113        usage: Option<&TokenBudgetUsage>,
114        cache_creation_input_tokens: u64,
115        cache_read_input_tokens: u64,
116        cache_write_input_tokens: u64,
117        input_tokens: u64,
118        output_tokens: u64,
119        thinking_tokens: u64,
120        context_management: ContextManagementTelemetry,
121    ) -> Self {
122        Self {
123            ts,
124            session_id: session_id.to_string(),
125            model: model.to_string(),
126            provider: provider.to_string(),
127            message_count,
128            cache_creation_input_tokens,
129            cache_read_input_tokens,
130            cache_write_input_tokens,
131            input_tokens,
132            output_tokens,
133            thinking_tokens,
134            system_tokens: usage.map(|u| u.system_tokens).unwrap_or(0),
135            summary_tokens: usage.map(|u| u.summary_tokens).unwrap_or(0),
136            window_tokens: usage.map(|u| u.window_tokens).unwrap_or(0),
137            total_tokens: usage.map(|u| u.total_tokens).unwrap_or(0),
138            max_context_tokens: usage.map(|u| u.max_context_tokens).unwrap_or(0),
139            budget_limit: usage.map(|u| u.budget_limit).unwrap_or(0),
140            prompt_cached_tool_outputs: usage.map(|u| u.prompt_cached_tool_outputs).unwrap_or(0),
141            prompt_cached_tool_tokens_saved: usage
142                .map(|u| u.prompt_cached_tool_tokens_saved)
143                .unwrap_or(0),
144            truncation_occurred: usage.map(|u| u.truncation_occurred).unwrap_or(false),
145            segments_removed: usage.map(|u| u.segments_removed).unwrap_or(0),
146            context_management_strategy: context_management.strategy,
147            model_context_epoch: context_management.model_context_epoch,
148            model_context_reset_reason: context_management.model_context_reset_reason,
149            retrieval_archive_event_count: context_management.retrieval_archive_event_count,
150            latest_retrieval_archive_event_id: context_management.latest_retrieval_archive_event_id,
151            latest_retrieval_archive_trigger_type: context_management
152                .latest_retrieval_archive_trigger_type,
153        }
154    }
155
156    /// Serialize to a single-line JSON string (no trailing newline; the storage
157    /// layer frames the line).
158    pub fn to_json_line(&self) -> Result<String, serde_json::Error> {
159        serde_json::to_string(self)
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn record_serializes_to_single_json_line_with_cache_creation() {
169        let usage = TokenBudgetUsage {
170            system_tokens: 5000,
171            summary_tokens: 2000,
172            window_tokens: 3000,
173            total_tokens: 10000,
174            max_context_tokens: 200_000,
175            budget_limit: 180_000,
176            truncation_occurred: false,
177            segments_removed: 0,
178            prompt_cached_tool_outputs: 1,
179            prompt_cached_tool_tokens_saved: 42,
180            thinking_tokens: 7,
181            cache_read_input_tokens: 12_000,
182            provider_prompt_usage: None,
183        };
184        let record = TokenUsageRecord::new(
185            "2026-06-15T00:00:00Z".to_string(),
186            "sess-1",
187            "claude-opus-4-8",
188            "anthropic",
189            24,
190            Some(&usage),
191            1500, // cache_creation — only present on the stream output
192            12_000,
193            0,
194            800, // input_tokens (fresh, non-cached)
195            300,
196            7,
197            ContextManagementTelemetry {
198                strategy: "retrieval_window".to_string(),
199                model_context_epoch: 4,
200                model_context_reset_reason: Some("compression".to_string()),
201                retrieval_archive_event_count: 2,
202                latest_retrieval_archive_event_id: Some("archive-event-2".to_string()),
203                latest_retrieval_archive_trigger_type: Some("auto".to_string()),
204            },
205        );
206        let line = record.to_json_line().expect("serializes");
207        assert!(!line.contains('\n'), "must be a single line");
208        assert!(line.contains("\"cache_creation_input_tokens\":1500"));
209        assert!(line.contains("\"cache_read_input_tokens\":12000"));
210        assert!(line.contains("\"cache_write_input_tokens\":0"));
211        assert!(line.contains("\"input_tokens\":800"));
212        assert!(line.contains("\"session_id\":\"sess-1\""));
213        assert!(line.contains("\"context_management_strategy\":\"retrieval_window\""));
214        assert!(line.contains("\"model_context_epoch\":4"));
215        assert!(line.contains("\"model_context_reset_reason\":\"compression\""));
216        assert!(line.contains("\"retrieval_archive_event_count\":2"));
217        assert!(line.contains("\"latest_retrieval_archive_event_id\":\"archive-event-2\""));
218        assert!(line.contains("\"latest_retrieval_archive_trigger_type\":\"auto\""));
219    }
220
221    #[test]
222    fn openai_flat_log_contract_keeps_fresh_cache_and_total_disjoint() {
223        let record = TokenUsageRecord::new(
224            "2026-07-29T00:00:00Z".to_string(),
225            "sess-openai",
226            "gpt-5",
227            "openai",
228            4,
229            None,
230            0,
231            768,
232            64,
233            232,
234            120,
235            20,
236            ContextManagementTelemetry::default(),
237        );
238
239        let prompt_total = record
240            .input_tokens
241            .saturating_add(record.cache_read_input_tokens)
242            .saturating_add(record.cache_creation_input_tokens);
243        let cache_hit_ratio = record.cache_read_input_tokens as f64 / prompt_total as f64;
244
245        assert_eq!(record.input_tokens, 232);
246        assert_eq!(prompt_total, 1000);
247        assert_eq!(record.cache_write_input_tokens, 64);
248        assert!((cache_hit_ratio - 0.768).abs() < f64::EPSILON);
249        assert_eq!(record.output_tokens, 120);
250        assert_eq!(record.thinking_tokens, 20);
251        assert_eq!(record.context_management_strategy, "summary");
252        assert_eq!(record.model_context_epoch, 0);
253        assert!(record.model_context_reset_reason.is_none());
254    }
255}
bamboo_engine/token_usage_log.rs

bamboo_engine/
token_usage_log.rs
