| 1 | //! Process-wide cost-accrual side-channel (#526). |
| 2 | //! |
| 3 | //! Background LLM calls outside the main turn-complete path |
| 4 | //! (compaction summaries, seam recompaction, cycle briefings) used |
| 5 | //! to drop their token usage on the floor — the dashboard's |
| 6 | //! session-cost only saw the parent turn's tokens, so a long |
| 7 | //! session that triggered compaction or cycle-restart under-reported |
| 8 | //! cost by however many tokens those background calls consumed. |
| 9 | //! |
| 10 | //! Mirrors the [`crate::retry_status`] pattern: background callers |
| 11 | //! call [`report`] after each `client.create_message`, the TUI |
| 12 | //! render loop calls [`drain`] every frame, and any drained amount |
| 13 | //! gets folded into `App::accrue_subagent_cost_estimate`. |
| 14 | //! |
| 15 | //! Why a side-channel and not a plumbed callback: the leaky callers |
| 16 | //! (`compaction.rs`, `seam_manager.rs`, `cycle_manager.rs`) are |
| 17 | //! engine-internal machinery without a direct handle to `App` or |
| 18 | //! the engine's event channel. A side-channel keeps the change |
| 19 | //! surface tiny — one new `report` line per call site — and any |
| 20 | //! future background caller (summarizers, retrieval helpers) gets |
| 21 | //! accrued for free without further plumbing. |
| 22 | |
| 23 | use std::sync::{Mutex, OnceLock}; |
| 24 | |
| 25 | use crate::models::Usage; |
| 26 | use crate::pricing::CostEstimate; |
| 27 | |
| 28 | static PENDING: OnceLock<Mutex<CostEstimate>> = OnceLock::new(); |
| 29 | |
| 30 | fn cell() -> &'static Mutex<CostEstimate> { |
| 31 | PENDING.get_or_init(|| Mutex::new(CostEstimate::default())) |
| 32 | } |
| 33 | |
| 34 | /// Background callers report their LLM usage here. Computes the |
| 35 | /// cost via [`crate::pricing::calculate_turn_cost_estimate_from_usage`] and |
| 36 | /// adds it to the pending pool. Cheap; takes a short-lived lock |
| 37 | /// and returns. No-op on models the pricing table doesn't know. |
| 38 | pub fn report(model: &str, usage: &Usage) { |
| 39 | let Some(cost) = crate::pricing::calculate_turn_cost_estimate_from_usage(model, usage) else { |
| 40 | return; |
| 41 | }; |
| 42 | if !cost.is_positive() { |
| 43 | return; |
| 44 | } |
| 45 | if let Ok(mut pending) = cell().lock() { |
| 46 | pending.usd += cost.usd; |
| 47 | pending.cny += cost.cny; |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | /// Drain the pending cost. Returns the accumulated amount and resets |
| 52 | /// the pool to zero. Called by the TUI render / event loop on each |
| 53 | /// frame; any non-zero result gets folded into `accrue_subagent_cost_estimate`. |
| 54 | pub fn drain() -> CostEstimate { |
| 55 | let Ok(mut pending) = cell().lock() else { |
| 56 | return CostEstimate::default(); |
| 57 | }; |
| 58 | std::mem::take(&mut *pending) |
| 59 | } |
| 60 | |
| 61 | /// Reset the pool to zero without consuming. Test-only helper for |
| 62 | /// suites that share the static and need to start from a known |
| 63 | /// state. Production code should always use [`drain`]. |
| 64 | #[cfg(test)] |
| 65 | pub fn reset_for_tests() { |
| 66 | if let Ok(mut pending) = cell().lock() { |
| 67 | *pending = CostEstimate::default(); |
| 68 | } |
| 69 | } |
| 70 | |
| 71 | #[cfg(test)] |
| 72 | mod tests { |
| 73 | use super::*; |
| 74 | |
| 75 | fn small_usage() -> Usage { |
| 76 | Usage { |
| 77 | input_tokens: 1_000, |
| 78 | output_tokens: 500, |
| 79 | ..Default::default() |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | /// Tests run in parallel and share the static — serialize the |
| 84 | /// ones that touch the pool through this mutex so concurrent |
| 85 | /// `report`/`drain` doesn't make assertions racy. |
| 86 | fn serial_lock() -> std::sync::MutexGuard<'static, ()> { |
| 87 | static M: OnceLock<Mutex<()>> = OnceLock::new(); |
| 88 | M.get_or_init(|| Mutex::new(())) |
| 89 | .lock() |
| 90 | .unwrap_or_else(|e| e.into_inner()) |
| 91 | } |
| 92 | |
| 93 | #[test] |
| 94 | fn report_adds_to_pool_and_drain_returns_then_resets() { |
| 95 | let _g = serial_lock(); |
| 96 | reset_for_tests(); |
| 97 | report("deepseek-v4-flash", &small_usage()); |
| 98 | let first = drain(); |
| 99 | assert!(first.usd > 0.0, "expected positive USD cost, got {first:?}"); |
| 100 | assert!(first.cny > 0.0, "expected positive CNY cost, got {first:?}"); |
| 101 | let second = drain(); |
| 102 | assert_eq!(second, CostEstimate::default(), "drain must zero the pool"); |
| 103 | } |
| 104 | |
| 105 | #[test] |
| 106 | fn report_skips_unknown_models() { |
| 107 | let _g = serial_lock(); |
| 108 | reset_for_tests(); |
| 109 | // NIM-hosted models intentionally have no DeepSeek pricing. |
| 110 | report("deepseek-ai/deepseek-v4-pro", &small_usage()); |
| 111 | assert_eq!(drain(), CostEstimate::default()); |
| 112 | } |
| 113 | |
| 114 | #[test] |
| 115 | fn report_accumulates_across_multiple_calls() { |
| 116 | let _g = serial_lock(); |
| 117 | reset_for_tests(); |
| 118 | report("deepseek-v4-flash", &small_usage()); |
| 119 | report("deepseek-v4-flash", &small_usage()); |
| 120 | let total = drain(); |
| 121 | // Two equal reports — total must be 2× a single report. |
| 122 | let single = crate::pricing::calculate_turn_cost_estimate_from_usage( |
| 123 | "deepseek-v4-flash", |
| 124 | &small_usage(), |
| 125 | ) |
| 126 | .unwrap(); |
| 127 | assert!((total.usd - 2.0 * single.usd).abs() < 1e-12); |
| 128 | assert!((total.cny - 2.0 * single.cny).abs() < 1e-12); |
| 129 | } |
| 130 | } |
| 131 |