1use std::collections::HashSet;
4use std::path::Path;
5
6use anyhow::{Context, Result, bail};
7use serde::{Deserialize, Serialize};
8
9use crate::core::bounce_tracker::BounceTracker;
10use crate::core::gain::model_pricing::{ModelPricing, PricingMatchKind};
11use crate::core::stats::{StatsStore, classify_command};
12
13const REPLAY_KIND: &str = "lean-ctx.quality-benchmark.v1";
14const MAX_REPLAY_BYTES: u64 = 16 * 1024 * 1024;
15const Z_95: f64 = 1.959_963_984_540_054;
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct ReplaySuite {
19 pub kind: String,
20 pub sessions: Vec<RecordedSession>,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct RecordedSession {
25 pub id: String,
26 pub model: String,
27 pub without_compression: RecordedArm,
28 pub with_compression: RecordedArm,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct RecordedArm {
33 pub success: bool,
34 pub turns: u64,
35 pub usage: TokenUsage,
36 #[serde(default)]
37 pub tool_calls: Vec<RecordedToolCall>,
38}
39
40#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
41pub struct TokenUsage {
42 pub new_input_tokens: u64,
43 pub cache_read_tokens: u64,
44 pub cache_write_tokens: u64,
45 pub output_tokens: u64,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct RecordedToolCall {
50 pub name: String,
51 #[serde(default)]
52 pub source: Option<String>,
53 #[serde(default)]
54 pub mode: Option<String>,
55 #[serde(default)]
56 pub original_tokens: u64,
57 #[serde(default)]
58 pub delivered_tokens: u64,
59}
60
61#[derive(Debug, Clone, Copy, PartialEq)]
62pub struct ConfidenceInterval {
63 pub low: f64,
64 pub high: f64,
65}
66
67#[derive(Debug, Clone)]
68pub struct ArmSummary {
69 pub successes: u64,
70 pub sessions: u64,
71 pub success_ci: ConfidenceInterval,
72 pub mean_turns: f64,
73 pub turns_ci: ConfidenceInterval,
74 pub tool_calls: u64,
75 pub expansions: u64,
76 pub bounce_ci: ConfidenceInterval,
77 pub total_cost_usd: f64,
78 pub mean_cost_usd: f64,
79 pub cost_ci: ConfidenceInterval,
80 pub compression_saved_tokens: u64,
81}
82
83#[derive(Debug, Clone)]
84pub struct ReplayReport {
85 pub without: ArmSummary,
86 pub with: ArmSummary,
87 pub success_delta: ConfidenceInterval,
88 pub mean_success_delta: f64,
89 pub turn_delta: ConfidenceInterval,
90 pub mean_turn_delta: f64,
91 pub cost_delta: ConfidenceInterval,
92 pub mean_cost_delta_usd: f64,
93 pub total_savings_usd: f64,
94}
95
96#[derive(Default)]
97struct ArmSamples {
98 successes: u64,
99 turns: Vec<f64>,
100 costs: Vec<f64>,
101 stats: StatsStore,
102 bounce: BounceTracker,
103}
104
105pub fn load_replay(path: &Path) -> Result<ReplaySuite> {
106 let metadata = std::fs::metadata(path)
107 .with_context(|| format!("reading replay metadata {}", path.display()))?;
108 if metadata.len() > MAX_REPLAY_BYTES {
109 bail!("replay exceeds {MAX_REPLAY_BYTES} byte limit");
110 }
111 let raw = std::fs::read_to_string(path)
112 .with_context(|| format!("reading replay {}", path.display()))?;
113 let suite: ReplaySuite =
114 serde_json::from_str(&raw).with_context(|| format!("parsing replay {}", path.display()))?;
115 validate(&suite)?;
116 Ok(suite)
117}
118
119fn validate(suite: &ReplaySuite) -> Result<()> {
120 if suite.kind != REPLAY_KIND {
121 bail!(
122 "unsupported replay kind {:?}; expected {REPLAY_KIND}",
123 suite.kind
124 );
125 }
126 if suite.sessions.is_empty() {
127 bail!("replay contains no sessions");
128 }
129 let mut ids = HashSet::new();
130 for session in &suite.sessions {
131 if session.id.trim().is_empty() || session.model.trim().is_empty() {
132 bail!("session id and model must be non-empty");
133 }
134 if !ids.insert(&session.id) {
135 bail!("duplicate session id {:?}", session.id);
136 }
137 for arm in [&session.without_compression, &session.with_compression] {
138 if arm.turns == 0 {
139 bail!("session {:?} has an arm with zero turns", session.id);
140 }
141 if arm
142 .tool_calls
143 .iter()
144 .any(|call| call.name.trim().is_empty())
145 {
146 bail!("session {:?} has a tool call without a name", session.id);
147 }
148 }
149 }
150 Ok(())
151}
152
153pub fn replay(suite: &ReplaySuite) -> Result<ReplayReport> {
154 validate(suite)?;
155 let pricing = ModelPricing::embedded();
156 let mut without = ArmSamples::default();
157 let mut with = ArmSamples::default();
158 let mut success_deltas = Vec::with_capacity(suite.sessions.len());
159 let mut turn_deltas = Vec::with_capacity(suite.sessions.len());
160 let mut cost_deltas = Vec::with_capacity(suite.sessions.len());
161
162 let mut sessions: Vec<_> = suite.sessions.iter().collect();
163 sessions.sort_unstable_by(|a, b| a.id.cmp(&b.id));
164 for session in sessions {
165 let quote = pricing.quote(Some(&session.model));
166 if quote.match_kind != PricingMatchKind::Exact {
167 bail!(
168 "session {:?} model {:?} has no exact embedded price",
169 session.id,
170 session.model
171 );
172 }
173 let baseline_cost = record_arm(&mut without, &session.without_compression, quote.cost);
174 let compressed_cost = record_arm(&mut with, &session.with_compression, quote.cost);
175 success_deltas.push(
176 f64::from(session.with_compression.success)
177 - f64::from(session.without_compression.success),
178 );
179 turn_deltas
180 .push(session.with_compression.turns as f64 - session.without_compression.turns as f64);
181 cost_deltas.push(compressed_cost - baseline_cost);
182 }
183
184 let without = summarize(&without);
185 let with = summarize(&with);
186 Ok(ReplayReport {
187 mean_success_delta: mean(&success_deltas),
188 success_delta: mean_ci(&success_deltas),
189 mean_turn_delta: mean(&turn_deltas),
190 turn_delta: mean_ci(&turn_deltas),
191 mean_cost_delta_usd: mean(&cost_deltas),
192 cost_delta: mean_ci(&cost_deltas),
193 total_savings_usd: without.total_cost_usd - with.total_cost_usd,
194 without,
195 with,
196 })
197}
198
199fn record_arm(
200 samples: &mut ArmSamples,
201 arm: &RecordedArm,
202 cost: crate::core::gain::model_pricing::ModelCost,
203) -> f64 {
204 samples.successes += u64::from(arm.success);
205 samples.turns.push(arm.turns as f64);
206 let usd = cost.estimate_usd(
207 arm.usage.new_input_tokens,
208 arm.usage.output_tokens,
209 arm.usage.cache_write_tokens,
210 arm.usage.cache_read_tokens,
211 );
212 samples.costs.push(usd);
213 for call in &arm.tool_calls {
214 samples.stats.total_commands = samples.stats.total_commands.saturating_add(1);
215 samples.stats.total_input_tokens = samples
216 .stats
217 .total_input_tokens
218 .saturating_add(call.original_tokens);
219 samples.stats.total_output_tokens = samples
220 .stats
221 .total_output_tokens
222 .saturating_add(call.delivered_tokens);
223 let entry = samples.stats.commands.entry(call.name.clone()).or_default();
224 entry.count = entry.count.saturating_add(1);
225 entry.input_tokens = entry.input_tokens.saturating_add(call.original_tokens);
226 entry.output_tokens = entry.output_tokens.saturating_add(call.delivered_tokens);
227 samples
228 .stats
229 .command_classes
230 .entry(call.name.clone())
231 .or_insert_with(|| classify_command(&call.name));
232
233 samples.bounce.next_seq();
234 if call.name == "ctx_expand" {
235 samples.bounce.record_expansion(
236 call.source.as_deref(),
237 usize::try_from(call.delivered_tokens).unwrap_or(usize::MAX),
238 );
239 }
240 }
241 usd
242}
243
244fn summarize(samples: &ArmSamples) -> ArmSummary {
245 let sessions = samples.turns.len() as u64;
246 let tool_calls = samples.stats.total_commands;
247 let expansions = samples.bounce.total_bounces();
248 let compression = samples.stats.compression_totals();
249 ArmSummary {
250 successes: samples.successes,
251 sessions,
252 success_ci: wilson_ci(samples.successes, sessions),
253 mean_turns: mean(&samples.turns),
254 turns_ci: mean_ci(&samples.turns),
255 tool_calls,
256 expansions,
257 bounce_ci: wilson_ci(expansions, tool_calls),
258 total_cost_usd: samples.costs.iter().sum(),
259 mean_cost_usd: mean(&samples.costs),
260 cost_ci: mean_ci(&samples.costs),
261 compression_saved_tokens: compression.saved_tokens(),
262 }
263}
264
265fn mean(values: &[f64]) -> f64 {
266 values.iter().sum::<f64>() / values.len() as f64
267}
268
269fn mean_ci(values: &[f64]) -> ConfidenceInterval {
270 let avg = mean(values);
271 if values.len() < 2 {
272 return ConfidenceInterval {
273 low: avg,
274 high: avg,
275 };
276 }
277 let variance =
278 values.iter().map(|v| (v - avg).powi(2)).sum::<f64>() / (values.len() - 1) as f64;
279 let margin = Z_95 * (variance / values.len() as f64).sqrt();
280 ConfidenceInterval {
281 low: avg - margin,
282 high: avg + margin,
283 }
284}
285
286fn wilson_ci(successes: u64, total: u64) -> ConfidenceInterval {
287 if total == 0 {
288 return ConfidenceInterval {
289 low: 0.0,
290 high: 0.0,
291 };
292 }
293 let n = total as f64;
294 let p = successes as f64 / n;
295 let z2 = Z_95 * Z_95;
296 let center = (p + z2 / (2.0 * n)) / (1.0 + z2 / n);
297 let margin = Z_95 * ((p * (1.0 - p) / n + z2 / (4.0 * n * n)).sqrt()) / (1.0 + z2 / n);
298 ConfidenceInterval {
299 low: (center - margin).max(0.0),
300 high: (center + margin).min(1.0),
301 }
302}
303
304pub fn format_markdown(report: &ReplayReport) -> String {
305 let arm = |name: &str, a: &ArmSummary| {
306 format!(
307 "| {name} | {}/{} ({:.2}% [{:.2}, {:.2}]) | {:.2} [{:.2}, {:.2}] | {}/{} ({:.2}% [{:.2}, {:.2}]) | ${:.6} (${:.6} [{:.6}, {:.6}]) | {} |",
308 a.successes,
309 a.sessions,
310 100.0 * a.successes as f64 / a.sessions as f64,
311 100.0 * a.success_ci.low,
312 100.0 * a.success_ci.high,
313 a.mean_turns,
314 a.turns_ci.low,
315 a.turns_ci.high,
316 a.expansions,
317 a.tool_calls,
318 if a.tool_calls == 0 {
319 0.0
320 } else {
321 100.0 * a.expansions as f64 / a.tool_calls as f64
322 },
323 100.0 * a.bounce_ci.low,
324 100.0 * a.bounce_ci.high,
325 a.total_cost_usd,
326 a.mean_cost_usd,
327 a.cost_ci.low,
328 a.cost_ci.high,
329 a.compression_saved_tokens,
330 )
331 };
332 let mut lines = vec![
333 "# lean-ctx Compression Quality A/B Benchmark".to_string(),
334 String::new(),
335 "95% confidence intervals; paired deltas are `with - without`. Costs use the embedded ModelPricing snapshot and all four billed token classes.".to_string(),
336 String::new(),
337 "| Arm | Task success | Turns/session | ctx_expand / tool calls | Total cost (mean/session) | Tool tokens saved |".to_string(),
338 "|---|---:|---:|---:|---:|---:|".to_string(),
339 arm("Without compression", &report.without),
340 arm("With compression", &report.with),
341 String::new(),
342 "## Paired impact".to_string(),
343 String::new(),
344 format!("- Success-rate delta: {:+.2} pp [{:+.2}, {:+.2}]", 100.0 * report.mean_success_delta, 100.0 * report.success_delta.low, 100.0 * report.success_delta.high),
345 format!("- Turn delta per task: {:+.3} [{:+.3}, {:+.3}]", report.mean_turn_delta, report.turn_delta.low, report.turn_delta.high),
346 format!("- Cost delta per task: {:+.6} USD [{:+.6}, {:+.6}]", report.mean_cost_delta_usd, report.cost_delta.low, report.cost_delta.high),
347 format!("- Total dollar savings: ${:.6}", report.total_savings_usd),
348 ];
349 lines.push(String::new());
350 lines.join("\n")
351}
352
353#[cfg(test)]
354mod tests;