Skip to main content

firstpass_proxy/
ope.rs

1//! Off-policy evaluation (OPE): direct-method replay and IPS/SNIPS for start-rung policies.
2//!
3//! # Direct method (ladder/threshold changes)
4//!
5//! Answers: "if I changed my ladder or serve threshold, what would my logged traffic have cost
6//! and how many failures would I have served?" β€” from the trace store, before enforcing anything.
7//!
8//! For each trace, walk the CANDIDATE ladder cheapest-first. At each rung, if the logged trace
9//! has an attempt for that exact model string, reuse its logged gate verdicts and cost. The first
10//! rung whose outcome would SERVE under the candidate rule ends the replay. If any candidate rung
11//! has no logged attempt the trace is UNEVALUABLE and excluded β€” never guessed.
12//!
13//! Coverage < 1.0 is the principal limitation: a candidate that adds models never logged cannot
14//! be evaluated. The report always surfaces coverage and n_correctness_known.
15//!
16//! # IPS / SNIPS for start-rung policies (requires exploration)
17//!
18//! Answers: "what would mean cost be if I always started at rung N?" β€” using importance-sampling
19//! on traffic logged under an epsilon-greedy policy that recorded propensities.
20//!
21//! Estimators (Horvitz-Thompson 1952 / Swaminathan-Joachims 2015):
22//! ```text
23//! wα΅’ = πŸ™[logged_start == N] / pα΅’          (importance weight)
24//! IPS  = (1/n) Ξ£α΅’ wα΅’ Β· costα΅’              (unbiased, higher variance)
25//! SNIPS = Ξ£(wα΅’Β·costα΅’) / Ξ£wα΅’              (self-normalised, lower variance)
26//! ESS  = (Ξ£wα΅’)Β² / Ξ£wα΅’Β²                   (effective sample size)
27//! ```
28//!
29//! Valid only when the candidate start rung N is in the logging policy's support (guaranteed
30//! under Ξ΅-greedy since p = Ξ΅/K > 0 for every rung). Traces without propensity are excluded
31//! and counted in `n_with_propensity`. The direct-method path remains for ladder-structure
32//! changes.
33//!
34//! # Doubly-robust (DR) estimator
35//!
36//! DR combines the direct-method baseline with an IPS residual correction:
37//! ```text
38//! DRα΅’ = DM(xα΅’, a) + wα΅’ Β· (rα΅’ βˆ’ DM(xα΅’, aα΅’))
39//! DR  = (1/n) Ξ£α΅’ DRα΅’
40//! ```
41//! where `DM(x, rung)` is the per-(task-kind, start-rung) empirical mean cost built from
42//! logged receipts. DR is unbiased when **either** the propensities **or** the reward model
43//! is correctly specified (double robustness). In practice, both are coarse approximations;
44//! DR tends to have lower variance than IPS while inheriting IPS's unbiasedness under
45//! correct propensities.
46
47use std::collections::HashMap;
48use std::path::Path;
49
50use firstpass_core::{Attempt, Config as RoutingConfig, DeferredVerdict, TaskKind, Trace, Verdict};
51
52use crate::calibrate::gate_score;
53use crate::store::{self, StoreError};
54
55// ── Candidate policy ──────────────────────────────────────────────────────────
56
57/// A candidate routing policy to evaluate against logged traffic.
58///
59/// Parsed from a standard `firstpass.toml` β€” the candidate is an ordinary config file; we
60/// extract the first route's ladder and `escalation.serve_threshold`.
61#[derive(Debug)]
62pub struct CandidatePolicy {
63    /// Model ladder, cheapest first (e.g. `"anthropic/claude-haiku-4-5"`).
64    pub ladder: Vec<String>,
65    /// Conformal serve threshold. `None` β†’ serve when verdict is `Pass`.
66    pub serve_threshold: Option<f64>,
67}
68
69impl CandidatePolicy {
70    /// Parse a candidate policy from a TOML string (standard firstpass config format).
71    ///
72    /// Uses the first `[[route]]` section's `ladder` and the top-level
73    /// `[escalation].serve_threshold`.
74    ///
75    /// # Errors
76    /// Returns a human-readable string if the TOML is invalid.
77    pub fn from_toml(toml: &str) -> Result<Self, String> {
78        let config = RoutingConfig::parse(toml).map_err(|e| e.to_string())?;
79        let ladder = config
80            .routes
81            .into_iter()
82            .next()
83            .map(|r| r.ladder)
84            .unwrap_or_default();
85        Ok(Self {
86            ladder,
87            serve_threshold: config.escalation.serve_threshold,
88        })
89    }
90}
91
92// ── Per-trace replay ──────────────────────────────────────────────────────────
93
94/// Whether a logged attempt would SERVE under the candidate policy.
95fn would_serve(attempt: &Attempt, policy: &CandidatePolicy) -> bool {
96    match policy.serve_threshold {
97        Some(t) => gate_score(&attempt.gates, attempt.verdict) >= t,
98        None => attempt.verdict == Verdict::Pass,
99    }
100}
101
102/// USD cost of one attempt: model call + all gate costs.
103fn attempt_cost(a: &Attempt) -> f64 {
104    a.cost_usd + a.gates.iter().map(|g| g.cost_usd).sum::<f64>()
105}
106
107enum ReplayResult {
108    /// A candidate rung had no logged attempt β€” the trace cannot be evaluated.
109    Unevaluable,
110    Evaluated {
111        /// Sum of attempt costs for every replayed rung (up to and including the serving one).
112        cost: f64,
113        /// Model string of the candidate rung that served, or `None` if all rungs exhausted.
114        served_model: Option<String>,
115        /// Model string actually served by the logging policy.
116        logged_served_model: Option<String>,
117    },
118}
119
120fn replay_trace(trace: &Trace, policy: &CandidatePolicy) -> ReplayResult {
121    // What the logging policy actually served β€” used for correctness attribution.
122    let logged_served_model = trace.final_.served_rung.and_then(|rung| {
123        trace
124            .attempts
125            .iter()
126            .find(|a| a.rung == rung)
127            .map(|a| a.model.clone())
128    });
129
130    let mut total_cost = 0.0f64;
131    let mut served_model: Option<String> = None;
132
133    for model in &policy.ladder {
134        let Some(attempt) = trace.attempts.iter().find(|a| &a.model == model) else {
135            return ReplayResult::Unevaluable;
136        };
137        total_cost += attempt_cost(attempt);
138        if would_serve(attempt, policy) {
139            served_model = Some(model.clone());
140            break;
141        }
142    }
143
144    ReplayResult::Evaluated {
145        cost: total_cost,
146        served_model,
147        logged_served_model,
148    }
149}
150
151// ── Bootstrap CIs ─────────────────────────────────────────────────────────────
152
153// ponytail: inline SplitMix64 β€” same pattern as conformal.rs tests; no new deps.
154// Ceiling: not a general RNG. Replace with rand if OPE ever needs sampling beyond CI bootstrap.
155fn splitmix64(state: &mut u64) -> u64 {
156    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
157    let mut z = *state;
158    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
159    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
160    z ^ (z >> 31)
161}
162
163#[inline]
164fn rand_usize(rng: &mut u64, n: usize) -> usize {
165    (splitmix64(rng) % n as u64) as usize
166}
167
168/// Bootstrap 95% CI for the mean of `values` (2.5/97.5 percentiles, deterministic seed).
169fn bootstrap_mean_ci(values: &[f64], n_resamples: usize, seed: u64) -> (f64, f64) {
170    if values.is_empty() {
171        return (0.0, 0.0);
172    }
173    let n = values.len();
174    let mut rng = seed;
175    let mut means: Vec<f64> = (0..n_resamples)
176        .map(|_| {
177            let s: f64 = (0..n).map(|_| values[rand_usize(&mut rng, n)]).sum();
178            s / n as f64
179        })
180        .collect();
181    means.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
182    let lo_idx = (n_resamples as f64 * 0.025) as usize;
183    let hi_idx = ((n_resamples as f64 * 0.975) as usize).min(n_resamples - 1);
184    (means[lo_idx], means[hi_idx])
185}
186
187/// Bootstrap 95% CI for a failure rate (2.5/97.5 percentiles, deterministic seed).
188fn bootstrap_failure_ci(correct: &[bool], n_resamples: usize, seed: u64) -> (f64, f64) {
189    if correct.is_empty() {
190        return (0.0, 0.0);
191    }
192    let n = correct.len();
193    let mut rng = seed;
194    let mut rates: Vec<f64> = (0..n_resamples)
195        .map(|_| {
196            let fails = (0..n).filter(|_| !correct[rand_usize(&mut rng, n)]).count();
197            fails as f64 / n as f64
198        })
199        .collect();
200    rates.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
201    let lo_idx = (n_resamples as f64 * 0.025) as usize;
202    let hi_idx = ((n_resamples as f64 * 0.975) as usize).min(n_resamples - 1);
203    (rates[lo_idx], rates[hi_idx])
204}
205
206// ── Report ────────────────────────────────────────────────────────────────────
207
208/// The result of evaluating a candidate policy against logged traffic.
209#[derive(Debug, Clone)]
210pub struct OpeReport {
211    /// Total traces in the store for this tenant.
212    pub n_traces: usize,
213    /// Traces for which replay used only logged attempts (coverage numerator).
214    pub n_evaluable: usize,
215    /// `n_evaluable / n_traces` (1.0 when n_traces == 0).
216    pub coverage: f64,
217    /// Mean candidate cost per request, over evaluable traces.
218    pub est_cost_per_request: f64,
219    /// Mean actual (logged) cost per request, over the same evaluable traces.
220    pub logged_cost_per_request: f64,
221    /// Estimated served-failure rate over evaluable traces with known correctness.
222    /// `None` if no evaluable trace has deferred feedback on the same served rung.
223    pub est_served_failure: Option<f64>,
224    /// Number of evaluable traces where correctness is attributable from deferred feedback.
225    pub n_correctness_known: usize,
226    /// Fraction of evaluable traces where candidate escalated past the first ladder rung.
227    pub escalation_rate: f64,
228    /// Bootstrap 95% CI (2.5/97.5 pct) for `est_cost_per_request`.
229    pub ci_cost: (f64, f64),
230    /// Bootstrap 95% CI for `est_served_failure`. `None` when `est_served_failure` is `None`.
231    pub ci_served_failure: Option<(f64, f64)>,
232}
233
234impl OpeReport {
235    /// Render the report as human-readable lines, mirroring `calibrate`'s report style.
236    /// Coverage and `n_correctness_known` are always prominent.
237    #[must_use]
238    pub fn render(&self) -> String {
239        let mut out = format!(
240            "traces: {n_traces}  evaluable: {n_evaluable}  coverage: {cov:.3}\n\
241             n_correctness_known: {n_known}\n\
242             est cost/request:    ${est:.6}  (logged: ${logged:.6})\n\
243             cost CI [2.5%, 97.5%]: [${clo:.6}, ${chi:.6}]\n\
244             escalation rate: {esc:.4}\n",
245            n_traces = self.n_traces,
246            n_evaluable = self.n_evaluable,
247            cov = self.coverage,
248            n_known = self.n_correctness_known,
249            est = self.est_cost_per_request,
250            logged = self.logged_cost_per_request,
251            clo = self.ci_cost.0,
252            chi = self.ci_cost.1,
253            esc = self.escalation_rate,
254        );
255        match self.est_served_failure {
256            Some(f) => {
257                let (lo, hi) = self.ci_served_failure.unwrap_or((f, f));
258                out.push_str(&format!(
259                    "est served-failure: {f:.4}  CI [{lo:.4}, {hi:.4}]\n"
260                ));
261            }
262            None => {
263                out.push_str(
264                    "est served-failure: n/a (no deferred feedback on same-rung evaluable traces)\n",
265                );
266            }
267        }
268        out.push_str(
269            "\nreplay of logged outcomes (direct method); \
270             rungs never logged are not guessed β€” see coverage.\n",
271        );
272        out
273    }
274}
275
276// ── Store-backed OPE ──────────────────────────────────────────────────────────
277
278/// Internal per-trace summary for aggregation and bootstrap resampling.
279struct EvalPoint {
280    candidate_cost: f64,
281    logged_cost: f64,
282    /// `Some(true)` = correct, `Some(false)` = failure, `None` = unknown.
283    correctness: Option<bool>,
284    escalated: bool,
285}
286
287fn build_report(n_traces: usize, points: Vec<EvalPoint>) -> OpeReport {
288    let n_evaluable = points.len();
289    // ponytail: coverage = 1.0 for empty store (no uncovered traces, not a zero).
290    let coverage = if n_traces == 0 {
291        1.0
292    } else {
293        n_evaluable as f64 / n_traces as f64
294    };
295
296    if points.is_empty() {
297        return OpeReport {
298            n_traces,
299            n_evaluable: 0,
300            coverage,
301            est_cost_per_request: 0.0,
302            logged_cost_per_request: 0.0,
303            est_served_failure: None,
304            n_correctness_known: 0,
305            escalation_rate: 0.0,
306            ci_cost: (0.0, 0.0),
307            ci_served_failure: None,
308        };
309    }
310
311    let est_cost = mean(&points, |p| p.candidate_cost);
312    let logged_cost = mean(&points, |p| p.logged_cost);
313    let escalation_rate = points.iter().filter(|p| p.escalated).count() as f64 / n_evaluable as f64;
314
315    let known: Vec<bool> = points.iter().filter_map(|p| p.correctness).collect();
316    let n_correctness_known = known.len();
317    let est_served_failure = if known.is_empty() {
318        None
319    } else {
320        Some(known.iter().filter(|&&c| !c).count() as f64 / known.len() as f64)
321    };
322
323    let costs: Vec<f64> = points.iter().map(|p| p.candidate_cost).collect();
324    // ponytail: fixed seeds (42/43) for determinism; 1000 resamples is the spec floor.
325    let ci_cost = bootstrap_mean_ci(&costs, 1000, 42);
326    let ci_served_failure = if known.is_empty() {
327        None
328    } else {
329        Some(bootstrap_failure_ci(&known, 1000, 43))
330    };
331
332    OpeReport {
333        n_traces,
334        n_evaluable,
335        coverage,
336        est_cost_per_request: est_cost,
337        logged_cost_per_request: logged_cost,
338        est_served_failure,
339        n_correctness_known,
340        escalation_rate,
341        ci_cost,
342        ci_served_failure,
343    }
344}
345
346fn mean(points: &[EvalPoint], f: impl Fn(&EvalPoint) -> f64) -> f64 {
347    if points.is_empty() {
348        return 0.0;
349    }
350    points.iter().map(f).sum::<f64>() / points.len() as f64
351}
352
353/// Run OPE from the trace store.
354///
355/// Missing or unreadable store is treated as zero traces (returns a zero-trace report and exits
356/// 0), matching the forgiving behaviour of `firstpass calibrate` and `firstpass trace`.
357///
358/// # Errors
359/// Returns [`StoreError`] if a stored trace's deferred verdicts cannot be read (a genuine
360/// database error on an existing record, not a missing store).
361pub fn ope_from_store(
362    db_path: impl AsRef<Path>,
363    tenant: &str,
364    policy: &CandidatePolicy,
365) -> Result<OpeReport, StoreError> {
366    let traces = store::load_tenant_traces(&db_path, tenant).unwrap_or_default();
367    let n_traces = traces.len();
368    let mut points: Vec<EvalPoint> = Vec::with_capacity(n_traces);
369
370    for trace in &traces {
371        let deferred = store::load_deferred(&db_path, &trace.trace_id.to_string())?;
372        let replay = replay_trace(trace, policy);
373        let ReplayResult::Evaluated {
374            cost,
375            served_model,
376            logged_served_model,
377        } = replay
378        else {
379            continue; // unevaluable: some candidate rung had no logged attempt
380        };
381
382        // Correctness is attributable only when candidate and logging policy served the same
383        // model AND deferred feedback exists. A different model means the deferred verdict graded
384        // a different output β€” we never impute correctness across outputs.
385        let correctness = match (&served_model, &logged_served_model) {
386            (Some(cm), Some(lm)) if cm == lm => deferred
387                .last()
388                .map(|dv: &DeferredVerdict| dv.verdict == Verdict::Pass),
389            _ => None,
390        };
391
392        let first_model = policy.ladder.first();
393        let escalated = match (&served_model, first_model) {
394            (Some(m), Some(first)) => m != first,
395            (None, Some(_)) => true, // all rungs exhausted without serving
396            _ => false,
397        };
398
399        points.push(EvalPoint {
400            candidate_cost: cost,
401            logged_cost: trace.final_.total_cost_usd,
402            correctness,
403            escalated,
404        });
405    }
406
407    Ok(build_report(n_traces, points))
408}
409
410// ── IPS / SNIPS start-rung estimator ─────────────────────────────────────────
411
412/// IPS, SNIPS, and doubly-robust (DR) estimates for evaluating a fixed candidate start-rung policy.
413///
414/// Valid only for traffic logged under a stochastic policy that recorded propensities via
415/// `[escalation.exploration]`. Traces without a `propensity` field are excluded and reported
416/// in `n_with_propensity`.
417///
418/// ponytail: per-context greedy identification is not exploited here (uniform-weight IPS
419/// suffices for population-level estimates).
420#[derive(Debug, Clone)]
421pub struct IpsReport {
422    /// Total traces in the store for this tenant.
423    pub n_traces: usize,
424    /// Traces with `propensity: Some(p > 0)` β€” the IPS population.
425    pub n_with_propensity: usize,
426    /// Candidate start rung being evaluated.
427    pub candidate_start_rung: u32,
428    /// IPS (Horvitz-Thompson) estimate of mean cost under "always start at rung N".
429    pub ips_cost: f64,
430    /// SNIPS (self-normalised IPS) estimate β€” lower variance, slightly biased.
431    pub snips_cost: f64,
432    /// Effective sample size ESS = (Ξ£wα΅’)Β² / Ξ£wα΅’Β².
433    pub ess: f64,
434    /// Bootstrap 95% CI (2.5/97.5 pct) for the IPS cost estimate.
435    pub ci_ips_cost: (f64, f64),
436    /// IPS estimate of served-failure rate (traces with deferred feedback on logged rung N).
437    pub ips_served_failure: Option<f64>,
438    /// SNIPS estimate of served-failure rate.
439    pub snips_served_failure: Option<f64>,
440    /// Traces contributing to the failure-rate IPS estimate.
441    pub n_correctness_known: usize,
442    /// Bootstrap 95% CI for the IPS served-failure estimate.
443    pub ci_ips_served_failure: Option<(f64, f64)>,
444    /// Doubly-robust (DR) cost estimate: DM baseline + IPS residual correction.
445    ///
446    /// `DRα΅’ = DM(xα΅’, N) + wα΅’ Β· (rα΅’ βˆ’ DM(xα΅’, aα΅’))`, mean over all propensity traces.
447    /// Unbiased when either propensities or the reward model is correctly specified.
448    pub dr_cost: f64,
449    /// Bootstrap 95% CI (2.5/97.5 pct) for the DR cost estimate.
450    pub ci_dr_cost: (f64, f64),
451}
452
453impl IpsReport {
454    /// Render the report as human-readable lines.
455    #[must_use]
456    pub fn render(&self) -> String {
457        let mut out = format!(
458            "traces: {n}  n_with_propensity: {nwp}  candidate_start_rung: {sr}\n\
459             IPS cost/request:   ${ips:.6}\n\
460             SNIPS cost/request: ${snips:.6}\n\
461             IPS cost CI [2.5%, 97.5%]: [${clo:.6}, ${chi:.6}]\n\
462             effective sample size (ESS): {ess:.1}\n\
463             n_correctness_known: {nck}\n",
464            n = self.n_traces,
465            nwp = self.n_with_propensity,
466            sr = self.candidate_start_rung,
467            ips = self.ips_cost,
468            snips = self.snips_cost,
469            clo = self.ci_ips_cost.0,
470            chi = self.ci_ips_cost.1,
471            ess = self.ess,
472            nck = self.n_correctness_known,
473        );
474        out.push_str(&format!(
475            "DR cost/request:    ${dr:.6}\n\
476             DR cost CI [2.5%, 97.5%]: [${drlo:.6}, ${drhi:.6}]\n",
477            dr = self.dr_cost,
478            drlo = self.ci_dr_cost.0,
479            drhi = self.ci_dr_cost.1,
480        ));
481        match self.ips_served_failure {
482            Some(f) => {
483                let sf_snips = self.snips_served_failure.unwrap_or(f);
484                let (lo, hi) = self.ci_ips_served_failure.unwrap_or((f, f));
485                out.push_str(&format!(
486                    "IPS served-failure: {f:.4}  SNIPS: {sf_snips:.4}  CI [{lo:.4}, {hi:.4}]\n"
487                ));
488            }
489            None => {
490                out.push_str(
491                    "IPS served-failure: n/a (no deferred feedback on matched-start traces)\n",
492                );
493            }
494        }
495        out.push_str(
496            "\nIPS/SNIPS/DR valid only for candidates over the same logged ladder with \
497             propensity-logged traffic. Traces without propensity excluded. \
498             DR reward model: per-(task-kind, rung) empirical mean β€” coarse; see ponytail comment. \
499             Direct-method replay remains for ladder/threshold changes.\n",
500        );
501        out
502    }
503}
504
505// ── DR reward model ───────────────────────────────────────────────────────────
506
507/// Per-(task-kind, start-rung) empirical mean cost, used as the DM baseline in DR.
508///
509/// Built from propensity-logged traces only (same population as IPS). Falls back to the
510/// per-rung mean, then the global mean, so `predict` never returns NaN even for (context,
511/// rung) pairs that were never explored.
512///
513/// ponytail: coarse empirical bucket mean β€” zero new deps, same reward definition as IPS.
514/// Ceiling: cannot extrapolate to (context, rung) pairs with zero logged observations; those
515/// fall back to the rung mean (or global mean), which may be biased if costs vary by context.
516/// Upgrade path: a richer feature regression (e.g. prompt-token bucket Γ— task kind Γ— rung)
517/// once the feature vector grows and sample sizes support it.
518struct DmModel {
519    /// Mean cost keyed by `(TaskKind, start_rung)`.
520    bucket: HashMap<(TaskKind, u32), f64>,
521    /// Per-rung mean cost (fallback when the context bucket is unobserved).
522    rung_mean: HashMap<u32, f64>,
523    /// Grand mean (ultimate fallback when even the rung has no observations).
524    global_mean: f64,
525}
526
527impl DmModel {
528    fn build(traces: &[Trace]) -> Self {
529        let mut bucket_acc: HashMap<(TaskKind, u32), (f64, u32)> = HashMap::new();
530        let mut rung_acc: HashMap<u32, (f64, u32)> = HashMap::new();
531        let mut global_sum = 0.0f64;
532        let mut global_n = 0u32;
533
534        for trace in traces {
535            if trace.policy.propensity.is_none() {
536                continue; // same filter as IPS β€” use only the propensity-logged population
537            }
538            let Some(first) = trace.attempts.first() else {
539                continue;
540            };
541            let rung = first.rung;
542            let cost = trace.final_.total_cost_usd;
543            let ctx = trace.request.features.task_kind;
544
545            let be = bucket_acc.entry((ctx, rung)).or_default();
546            be.0 += cost;
547            be.1 += 1;
548
549            let re = rung_acc.entry(rung).or_default();
550            re.0 += cost;
551            re.1 += 1;
552
553            global_sum += cost;
554            global_n += 1;
555        }
556
557        let bucket = bucket_acc
558            .into_iter()
559            .map(|(k, (s, n))| (k, s / n as f64))
560            .collect();
561        let rung_mean = rung_acc
562            .into_iter()
563            .map(|(k, (s, n))| (k, s / n as f64))
564            .collect();
565        let global_mean = if global_n > 0 {
566            global_sum / global_n as f64
567        } else {
568            0.0
569        };
570
571        Self {
572            bucket,
573            rung_mean,
574            global_mean,
575        }
576    }
577
578    /// Predicted cost for `(task_kind, rung)`, with graceful fallback.
579    fn predict(&self, task_kind: TaskKind, rung: u32) -> f64 {
580        self.bucket
581            .get(&(task_kind, rung))
582            .copied()
583            .or_else(|| self.rung_mean.get(&rung).copied())
584            .unwrap_or(self.global_mean)
585    }
586}
587
588/// Run IPS/SNIPS OPE from the trace store for a fixed candidate start-rung policy.
589///
590/// Traces without `propensity` are excluded from IPS and counted in `n_with_propensity`.
591/// Missing or unreadable store is treated as zero traces (same forgiving behaviour as
592/// [`ope_from_store`]).
593///
594/// # Errors
595/// Returns [`StoreError`] on a genuine database read error on an existing record.
596pub fn ips_from_store(
597    db_path: impl AsRef<Path>,
598    tenant: &str,
599    candidate_start_rung: u32,
600) -> Result<IpsReport, StoreError> {
601    let traces = store::load_tenant_traces(&db_path, tenant).unwrap_or_default();
602    let n_traces = traces.len();
603
604    struct IpsPoint {
605        weight: f64,
606        cost: f64,
607        /// `Some(true)` correct, `Some(false)` failure, `None` unknown.
608        correctness: Option<bool>,
609        // Needed for DR DM-term lookups.
610        task_kind: TaskKind,
611        /// Logged start rung; 0 when the trace has no attempts (w=0 so correction is 0).
612        logged_rung: u32,
613    }
614
615    let mut points: Vec<IpsPoint> = Vec::with_capacity(n_traces);
616    let mut n_with_propensity = 0usize;
617
618    for trace in &traces {
619        let Some(p) = trace.policy.propensity else {
620            continue; // no propensity β†’ excluded from IPS
621        };
622        if p <= 0.0 {
623            continue; // defensive: zero propensity β†’ undefined ratio
624        }
625        n_with_propensity += 1;
626
627        let logged_start = trace.attempts.first().map(|a| a.rung);
628        let indicator = f64::from(logged_start == Some(candidate_start_rung));
629        let w = indicator / p;
630
631        // Correctness from deferred feedback, but only for traces that matched the candidate
632        // start rung (w > 0). A deferred verdict on a different rung doesn't grade rung N's output.
633        let correctness = if w > 0.0 {
634            let deferred = store::load_deferred(&db_path, &trace.trace_id.to_string())?;
635            deferred
636                .last()
637                .map(|dv: &DeferredVerdict| dv.verdict == Verdict::Pass)
638        } else {
639            None
640        };
641
642        points.push(IpsPoint {
643            weight: w,
644            cost: trace.final_.total_cost_usd,
645            correctness,
646            task_kind: trace.request.features.task_kind,
647            logged_rung: logged_start.unwrap_or(0),
648        });
649    }
650
651    let n = n_with_propensity as f64;
652    let sum_w: f64 = points.iter().map(|p| p.weight).sum();
653    let sum_w2: f64 = points.iter().map(|p| p.weight * p.weight).sum();
654    let sum_wc: f64 = points.iter().map(|p| p.weight * p.cost).sum();
655
656    let ips_cost = if n > 0.0 { sum_wc / n } else { 0.0 };
657    let snips_cost = if sum_w > 0.0 { sum_wc / sum_w } else { 0.0 };
658    let ess = if sum_w2 > 0.0 {
659        sum_w * sum_w / sum_w2
660    } else {
661        0.0
662    };
663
664    // Bootstrap CI for IPS cost: resample the per-trace wα΅’Β·costα΅’ values.
665    let wc_values: Vec<f64> = points.iter().map(|p| p.weight * p.cost).collect();
666    let ci_ips_cost = bootstrap_mean_ci(&wc_values, 1000, 42);
667
668    // Failure-rate IPS over traces with deferred feedback at the matched start rung (w > 0).
669    let known: Vec<(f64, bool)> = points
670        .iter()
671        .filter_map(|p| p.correctness.map(|c| (p.weight, c)))
672        .collect();
673    let n_correctness_known = known.len();
674
675    let (ips_served_failure, snips_served_failure, ci_ips_served_failure) = if known.is_empty() {
676        (None, None, None)
677    } else {
678        let n_known = known.len() as f64;
679        let sum_wf: f64 = known.iter().filter(|(_, c)| !c).map(|(w, _)| w).sum();
680        let sum_wk: f64 = known.iter().map(|(w, _)| w).sum();
681        let ips_f = sum_wf / n_known;
682        let snips_f = if sum_wk > 0.0 { sum_wf / sum_wk } else { 0.0 };
683        let wf_vals: Vec<f64> = known
684            .iter()
685            .map(|(w, c)| if !c { *w } else { 0.0 })
686            .collect();
687        let ci = bootstrap_mean_ci(&wf_vals, 1000, 43);
688        (Some(ips_f), Some(snips_f), Some(ci))
689    };
690
691    // ── DR estimator ──────────────────────────────────────────────────────────
692    // DRα΅’ = DM(xα΅’, N) + wα΅’ Β· (rα΅’ βˆ’ DM(xα΅’, aα΅’))
693    // Averaged over ALL propensity-positive traces (including w=0 ones, which contribute
694    // only the DM baseline term). This is what reduces DR's variance relative to IPS.
695    let dm = DmModel::build(&traces);
696    let dr_vals: Vec<f64> = points
697        .iter()
698        .map(|pt| {
699            let dm_cand = dm.predict(pt.task_kind, candidate_start_rung);
700            let dm_logged = dm.predict(pt.task_kind, pt.logged_rung);
701            dm_cand + pt.weight * (pt.cost - dm_logged)
702        })
703        .collect();
704    let dr_cost = if dr_vals.is_empty() {
705        0.0
706    } else {
707        dr_vals.iter().sum::<f64>() / dr_vals.len() as f64
708    };
709    // ponytail: seed 44 β€” unique from IPS seeds 42/43 for determinism.
710    let ci_dr_cost = bootstrap_mean_ci(&dr_vals, 1000, 44);
711
712    Ok(IpsReport {
713        n_traces,
714        n_with_propensity,
715        candidate_start_rung,
716        ips_cost,
717        snips_cost,
718        ess,
719        ci_ips_cost,
720        ips_served_failure,
721        snips_served_failure,
722        n_correctness_known,
723        ci_ips_served_failure,
724        dr_cost,
725        ci_dr_cost,
726    })
727}
728
729// ── Tests ─────────────────────────────────────────────────────────────────────
730
731#[cfg(test)]
732mod tests {
733    use firstpass_core::{
734        Features, FinalOutcome, GENESIS_HASH, GateResult, Mode, PolicyRef, RequestInfo, Score,
735        ServedFrom, TaskKind, Verdict,
736    };
737
738    use super::*;
739    use crate::store;
740
741    /// Minimal trace with one attempt at `rung`, model `model_str`, verdict and score set by
742    /// `pass_score` (>= 0.5 = Pass). Mirrors calibrate.rs's `trace_with_score`.
743    fn make_trace(tenant: &str, rung: u32, model: &str, pass_score: f64, cost_usd: f64) -> Trace {
744        let verdict = if pass_score >= 0.5 {
745            Verdict::Pass
746        } else {
747            Verdict::Fail
748        };
749        let attempt = firstpass_core::Attempt {
750            rung,
751            model: model.to_owned(),
752            provider: "anthropic".to_owned(),
753            in_tokens: 10,
754            out_tokens: 5,
755            cost_usd,
756            latency_ms: 12,
757            gates: vec![GateResult {
758                gate_id: "gate@v1".to_owned(),
759                verdict,
760                score: Some(Score::clamped(pass_score)),
761                cost_usd: 0.0,
762                ms: 10,
763                reason: None,
764                evidence_ref: None,
765            }],
766            verdict,
767        };
768        let mut trace = Trace {
769            trace_id: uuid::Uuid::now_v7(),
770            prev_hash: GENESIS_HASH.to_owned(),
771            tenant_id: tenant.to_owned(),
772            session_id: "s1".to_owned(),
773            ts: jiff::Timestamp::now(),
774            mode: Mode::Enforce,
775            policy: PolicyRef {
776                id: "test@v0".to_owned(),
777                explore: false,
778                propensity: None,
779                mode_profile: None,
780            },
781            request: RequestInfo {
782                api: "anthropic.messages".to_owned(),
783                prompt_hash: "deadbeef".to_owned(),
784                features: Features::new(TaskKind::Other),
785            },
786            attempts: vec![attempt],
787            deferred: Vec::new(),
788            final_: FinalOutcome {
789                served_rung: Some(rung),
790                served_from: ServedFrom::Attempt,
791                total_cost_usd: cost_usd,
792                gate_cost_usd: 0.0,
793                total_latency_ms: 12,
794                escalations: 0,
795                counterfactual_baseline_usd: cost_usd,
796                savings_usd: 0.0,
797            },
798            probe: None,
799            rollout: None,
800            shadow: None,
801            predicted_pass: None,
802            elastic: None,
803        };
804        trace.recompute_savings();
805        trace
806    }
807
808    /// Two-attempt trace: haiku (rung 0) fails, sonnet (rung 1) passes.
809    fn make_escalated_trace(tenant: &str, haiku_cost: f64, sonnet_cost: f64) -> Trace {
810        let haiku = firstpass_core::Attempt {
811            rung: 0,
812            model: "haiku".to_owned(),
813            provider: "anthropic".to_owned(),
814            in_tokens: 10,
815            out_tokens: 5,
816            cost_usd: haiku_cost,
817            latency_ms: 10,
818            gates: vec![GateResult {
819                gate_id: "g".to_owned(),
820                verdict: Verdict::Fail,
821                score: Some(Score::clamped(0.3)),
822                cost_usd: 0.0,
823                ms: 5,
824                reason: None,
825                evidence_ref: None,
826            }],
827            verdict: Verdict::Fail,
828        };
829        let sonnet = firstpass_core::Attempt {
830            rung: 1,
831            model: "sonnet".to_owned(),
832            provider: "anthropic".to_owned(),
833            in_tokens: 10,
834            out_tokens: 5,
835            cost_usd: sonnet_cost,
836            latency_ms: 20,
837            gates: vec![GateResult {
838                gate_id: "g".to_owned(),
839                verdict: Verdict::Pass,
840                score: Some(Score::clamped(0.9)),
841                cost_usd: 0.0,
842                ms: 5,
843                reason: None,
844                evidence_ref: None,
845            }],
846            verdict: Verdict::Pass,
847        };
848        let total = haiku_cost + sonnet_cost;
849        let mut trace = Trace {
850            trace_id: uuid::Uuid::now_v7(),
851            prev_hash: GENESIS_HASH.to_owned(),
852            tenant_id: tenant.to_owned(),
853            session_id: "s2".to_owned(),
854            ts: jiff::Timestamp::now(),
855            mode: Mode::Enforce,
856            policy: PolicyRef {
857                id: "test@v0".to_owned(),
858                explore: false,
859                propensity: None,
860                mode_profile: None,
861            },
862            request: RequestInfo {
863                api: "anthropic.messages".to_owned(),
864                prompt_hash: "beef".to_owned(),
865                features: Features::new(TaskKind::Other),
866            },
867            attempts: vec![haiku, sonnet],
868            deferred: Vec::new(),
869            final_: FinalOutcome {
870                served_rung: Some(1),
871                served_from: ServedFrom::Attempt,
872                total_cost_usd: total,
873                gate_cost_usd: 0.0,
874                total_latency_ms: 30,
875                escalations: 1,
876                counterfactual_baseline_usd: total,
877                savings_usd: 0.0,
878            },
879            probe: None,
880            rollout: None,
881            shadow: None,
882            predicted_pass: None,
883            elastic: None,
884        };
885        trace.recompute_savings();
886        trace
887    }
888
889    fn deferred_pass(gate: &str) -> firstpass_core::DeferredVerdict {
890        firstpass_core::DeferredVerdict {
891            gate_id: gate.to_owned(),
892            verdict: Verdict::Pass,
893            score: None,
894            reported_at: jiff::Timestamp::now(),
895            reporter: "test".to_owned(),
896        }
897    }
898
899    fn deferred_fail(gate: &str) -> firstpass_core::DeferredVerdict {
900        firstpass_core::DeferredVerdict {
901            gate_id: gate.to_owned(),
902            verdict: Verdict::Fail,
903            score: None,
904            reported_at: jiff::Timestamp::now(),
905            reporter: "test".to_owned(),
906        }
907    }
908
909    fn tmp_db() -> std::path::PathBuf {
910        std::env::temp_dir().join(format!("fp-ope-{}.db", uuid::Uuid::now_v7()))
911    }
912
913    // ── 1. Pure replay: candidate == logged ladder ────────────────────────────
914
915    #[tokio::test]
916    async fn candidate_equals_logged_matches_exactly() {
917        let db = tmp_db();
918        let (tx, handle) = store::open(&db).unwrap();
919
920        // 10 traces, haiku passes at cost 0.001 each.
921        let mut ids = Vec::new();
922        for _ in 0..10 {
923            let t = make_trace("tenant-a", 0, "haiku", 0.8, 0.001);
924            ids.push(t.trace_id.to_string());
925            tx.try_send(t).unwrap();
926        }
927        drop(tx);
928        handle.await.unwrap();
929
930        for id in &ids {
931            store::append_deferred(&db, id, &deferred_pass("out")).unwrap();
932        }
933
934        let policy = CandidatePolicy {
935            ladder: vec!["haiku".to_owned()],
936            serve_threshold: None,
937        };
938        let report = ope_from_store(&db, "tenant-a", &policy).unwrap();
939
940        assert_eq!(report.n_traces, 10);
941        assert_eq!(report.n_evaluable, 10);
942        assert!((report.coverage - 1.0).abs() < 1e-9);
943        // Cost must match exactly: candidate replays the same attempt.
944        assert!((report.est_cost_per_request - 0.001).abs() < 1e-9);
945        assert!((report.logged_cost_per_request - 0.001).abs() < 1e-9);
946        // All served same rung as logged, all deferred = Pass -> failure = 0.
947        assert_eq!(report.n_correctness_known, 10);
948        assert!((report.est_served_failure.unwrap() - 0.0).abs() < 1e-9);
949        // No escalation: haiku always served first.
950        assert!((report.escalation_rate - 0.0).abs() < 1e-9);
951
952        let _ = std::fs::remove_file(&db);
953    }
954
955    // ── 2. Cheaper candidate (first rung only, drop sonnet) ────────────────────
956
957    #[tokio::test]
958    async fn cheaper_candidate_reduces_cost_and_escalation() {
959        let db = tmp_db();
960        let (tx, handle) = store::open(&db).unwrap();
961
962        // 5 traces logged: haiku (cost 0.001) fails -> sonnet (cost 0.01) passes.
963        // Logged total = 0.011 each.
964        for _ in 0..5 {
965            tx.try_send(make_escalated_trace("t", 0.001, 0.01)).unwrap();
966        }
967        drop(tx);
968        handle.await.unwrap();
969
970        // Candidate: ladder = ["haiku"] only.
971        // Replay: haiku fails -> all candidate rungs exhausted. Cost = 0.001.
972        // Logged served sonnet; candidate served nothing -> correctness UNKNOWN.
973        let policy = CandidatePolicy {
974            ladder: vec!["haiku".to_owned()],
975            serve_threshold: None,
976        };
977        let report = ope_from_store(&db, "t", &policy).unwrap();
978
979        assert_eq!(report.n_evaluable, 5);
980        assert!((report.coverage - 1.0).abs() < 1e-9);
981        // Candidate cost = 0.001 (haiku only); logged = 0.011.
982        assert!(
983            (report.est_cost_per_request - 0.001).abs() < 1e-9,
984            "got {}",
985            report.est_cost_per_request
986        );
987        assert!((report.logged_cost_per_request - 0.011).abs() < 1e-9);
988        // Served nothing vs logged sonnet -> UNKNOWN for all.
989        assert_eq!(report.n_correctness_known, 0);
990        assert!(report.est_served_failure.is_none());
991        // Escalated: candidate exhausted past first rung (haiku didn't serve).
992        assert!((report.escalation_rate - 1.0).abs() < 1e-9);
993
994        let _ = std::fs::remove_file(&db);
995    }
996
997    // ── 3. Candidate with an unlogged model -> unevaluable ────────────────────
998
999    #[tokio::test]
1000    async fn unlogged_model_makes_trace_unevaluable() {
1001        let db = tmp_db();
1002        let (tx, handle) = store::open(&db).unwrap();
1003
1004        // 3 traces have haiku attempts; 3 traces have sonnet attempts.
1005        for _ in 0..3 {
1006            tx.try_send(make_trace("t", 0, "haiku", 0.8, 0.001))
1007                .unwrap();
1008        }
1009        for _ in 0..3 {
1010            tx.try_send(make_trace("t", 0, "sonnet", 0.8, 0.01))
1011                .unwrap();
1012        }
1013        drop(tx);
1014        handle.await.unwrap();
1015
1016        // Candidate ladder: ["newmodel", "haiku"] β€” "newmodel" is never logged.
1017        // Every trace is unevaluable (first candidate rung has no logged attempt).
1018        let policy = CandidatePolicy {
1019            ladder: vec!["newmodel".to_owned(), "haiku".to_owned()],
1020            serve_threshold: None,
1021        };
1022        let report = ope_from_store(&db, "t", &policy).unwrap();
1023
1024        assert_eq!(report.n_traces, 6);
1025        assert_eq!(report.n_evaluable, 0);
1026        assert!((report.coverage - 0.0).abs() < 1e-9);
1027        assert!(report.est_served_failure.is_none());
1028
1029        let _ = std::fs::remove_file(&db);
1030    }
1031
1032    // ── 4. Different rung served β†’ correctness UNKNOWN ────────────────────────
1033
1034    #[tokio::test]
1035    async fn different_rung_served_correctness_unknown() {
1036        let db = tmp_db();
1037        let (tx, handle) = store::open(&db).unwrap();
1038
1039        // Logged: haiku at rung 0 fails, sonnet at rung 1 passes. Logged serves sonnet.
1040        let t = make_escalated_trace("t", 0.001, 0.01);
1041        let tid = t.trace_id.to_string();
1042        tx.try_send(t).unwrap();
1043        drop(tx);
1044        handle.await.unwrap();
1045
1046        // Deferred feedback graded the sonnet output.
1047        store::append_deferred(&db, &tid, &deferred_pass("out")).unwrap();
1048
1049        // Candidate: ladder = ["haiku", "sonnet"] with a very low threshold -> haiku serves.
1050        // Candidate serves haiku, logged served sonnet -> DIFFERENT rung -> UNKNOWN.
1051        let policy = CandidatePolicy {
1052            ladder: vec!["haiku".to_owned(), "sonnet".to_owned()],
1053            serve_threshold: Some(0.1), // haiku score=0.3 >= 0.1, so haiku would serve
1054        };
1055        let report = ope_from_store(&db, "t", &policy).unwrap();
1056
1057        assert_eq!(report.n_evaluable, 1);
1058        assert_eq!(report.n_correctness_known, 0, "different rung => UNKNOWN");
1059        assert!(report.est_served_failure.is_none());
1060        // Candidate served on haiku (cost 0.001), logged served on haiku+sonnet (0.011).
1061        assert!((report.est_cost_per_request - 0.001).abs() < 1e-9);
1062        // No escalation: haiku served first.
1063        assert!((report.escalation_rate - 0.0).abs() < 1e-9);
1064
1065        let _ = std::fs::remove_file(&db);
1066    }
1067
1068    // ── 5. Bootstrap CI: deterministic, contains point estimate ───────────────
1069
1070    #[tokio::test]
1071    async fn bootstrap_ci_deterministic_and_sane() {
1072        let db = tmp_db();
1073        let (tx, handle) = store::open(&db).unwrap();
1074
1075        // 30 traces with varying costs: alternating 0.001 and 0.002. Mean = 0.0015.
1076        let mut ids = Vec::new();
1077        for i in 0..30u32 {
1078            let cost = if i % 2 == 0 { 0.001 } else { 0.002 };
1079            let t = make_trace("t", 0, "m", 0.9, cost);
1080            ids.push(t.trace_id.to_string());
1081            tx.try_send(t).unwrap();
1082        }
1083        drop(tx);
1084        handle.await.unwrap();
1085
1086        // Half pass, half fail (deferred).
1087        for (i, id) in ids.iter().enumerate() {
1088            let dv = if i < 15 {
1089                deferred_pass("o")
1090            } else {
1091                deferred_fail("o")
1092            };
1093            store::append_deferred(&db, id, &dv).unwrap();
1094        }
1095
1096        let policy = CandidatePolicy {
1097            ladder: vec!["m".to_owned()],
1098            serve_threshold: None,
1099        };
1100        let r1 = ope_from_store(&db, "t", &policy).unwrap();
1101        let r2 = ope_from_store(&db, "t", &policy).unwrap();
1102
1103        // Deterministic: same CI both calls.
1104        assert_eq!(r1.ci_cost, r2.ci_cost, "CI must be deterministic");
1105        assert_eq!(r1.ci_served_failure, r2.ci_served_failure);
1106
1107        // CI ordering: lo <= point estimate <= hi.
1108        let (lo, hi) = r1.ci_cost;
1109        assert!(
1110            lo <= r1.est_cost_per_request + 1e-9,
1111            "CI lo {lo} > est {}",
1112            r1.est_cost_per_request
1113        );
1114        assert!(
1115            hi >= r1.est_cost_per_request - 1e-9,
1116            "CI hi {hi} < est {}",
1117            r1.est_cost_per_request
1118        );
1119        assert!(lo <= hi, "CI must be ordered");
1120
1121        if let Some((flo, fhi)) = r1.ci_served_failure {
1122            let f = r1.est_served_failure.unwrap();
1123            assert!(flo <= f + 1e-9, "failure CI lo {flo} > est {f}");
1124            assert!(fhi >= f - 1e-9, "failure CI hi {fhi} < est {f}");
1125            assert!(flo <= fhi);
1126        }
1127
1128        let _ = std::fs::remove_file(&db);
1129    }
1130
1131    // ── 6. Empty store β†’ zero-trace report, exit-0 path ──────────────────────
1132
1133    #[test]
1134    fn empty_store_returns_zero_trace_report() {
1135        // Non-existent db: load_tenant_traces returns empty, should give a valid zero report.
1136        let policy = CandidatePolicy {
1137            ladder: vec!["m".to_owned()],
1138            serve_threshold: None,
1139        };
1140        let db = std::path::Path::new("/nonexistent/fp-ope-empty.db");
1141        let report = ope_from_store(db, "t", &policy).unwrap();
1142        assert_eq!(report.n_traces, 0);
1143        assert_eq!(report.n_evaluable, 0);
1144        assert!((report.coverage - 1.0).abs() < 1e-9);
1145        assert!(report.est_served_failure.is_none());
1146    }
1147
1148    // ── 7. Same-rung with deferred: correctness flows through ─────────────────
1149
1150    #[tokio::test]
1151    async fn same_rung_deferred_correctness_attributed() {
1152        let db = tmp_db();
1153        let (tx, handle) = store::open(&db).unwrap();
1154
1155        let t_pass = make_trace("t", 0, "m", 0.9, 0.001);
1156        let t_fail = make_trace("t", 0, "m", 0.9, 0.001);
1157        let (id_pass, id_fail) = (t_pass.trace_id.to_string(), t_fail.trace_id.to_string());
1158        tx.try_send(t_pass).unwrap();
1159        tx.try_send(t_fail).unwrap();
1160        drop(tx);
1161        handle.await.unwrap();
1162
1163        store::append_deferred(&db, &id_pass, &deferred_pass("out")).unwrap();
1164        store::append_deferred(&db, &id_fail, &deferred_fail("out")).unwrap();
1165
1166        let policy = CandidatePolicy {
1167            ladder: vec!["m".to_owned()],
1168            serve_threshold: None,
1169        };
1170        let r = ope_from_store(&db, "t", &policy).unwrap();
1171
1172        assert_eq!(r.n_evaluable, 2);
1173        assert_eq!(r.n_correctness_known, 2);
1174        // 1 failure out of 2 known = 0.5.
1175        assert!((r.est_served_failure.unwrap() - 0.5).abs() < 1e-9);
1176
1177        let _ = std::fs::remove_file(&db);
1178    }
1179
1180    // ── 8. Serve threshold raises cost (candidate needs more escalations) ──────
1181
1182    #[tokio::test]
1183    async fn high_threshold_forces_escalation_and_higher_cost() {
1184        let db = tmp_db();
1185        let (tx, handle) = store::open(&db).unwrap();
1186
1187        // Logged: haiku (score=0.7, cost 0.001) passes under default rule (Pass verdict).
1188        // Candidate uses serve_threshold=0.8 -> haiku score 0.7 < 0.8 -> fail -> escalate.
1189        // Trace has no sonnet attempt -> trace is UNEVALUABLE (candidate needs sonnet, not logged).
1190        let t = make_trace("t", 0, "haiku", 0.7, 0.001); // score 0.7, verdict Pass
1191        tx.try_send(t).unwrap();
1192        drop(tx);
1193        handle.await.unwrap();
1194
1195        let policy = CandidatePolicy {
1196            ladder: vec!["haiku".to_owned()],
1197            serve_threshold: Some(0.8), // haiku score 0.7 < 0.8 -> no serve
1198        };
1199        let r = ope_from_store(&db, "t", &policy).unwrap();
1200
1201        // Haiku is logged, candidate walks it (score 0.7 < 0.8, no serve), exhausts all
1202        // candidate rungs β€” evaluable (all candidate rungs had logged data), no serve.
1203        assert_eq!(r.n_evaluable, 1);
1204        assert!((r.est_cost_per_request - 0.001).abs() < 1e-9); // haiku cost still counted
1205        assert_eq!(r.n_correctness_known, 0); // candidate served nothing vs logged haiku
1206        // Escalated: exhausted past first candidate rung.
1207        assert!((r.escalation_rate - 1.0).abs() < 1e-9);
1208
1209        let _ = std::fs::remove_file(&db);
1210    }
1211
1212    // ── 9. CandidatePolicy::from_toml parses ladder and threshold ─────────────
1213
1214    #[test]
1215    fn from_toml_extracts_first_route_and_threshold() {
1216        let toml = r#"
1217[[route]]
1218match = {}
1219mode = "enforce"
1220ladder = ["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"]
1221
1222[escalation]
1223serve_threshold = 0.75
1224"#;
1225        let p = CandidatePolicy::from_toml(toml).unwrap();
1226        assert_eq!(
1227            p.ladder,
1228            ["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"]
1229        );
1230        assert!((p.serve_threshold.unwrap() - 0.75).abs() < 1e-9);
1231    }
1232
1233    #[test]
1234    fn from_toml_no_threshold_is_none() {
1235        let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"m\"]\n";
1236        let p = CandidatePolicy::from_toml(toml).unwrap();
1237        assert!(p.serve_threshold.is_none());
1238    }
1239
1240    // ── IPS / SNIPS tests ─────────────────────────────────────────────────────
1241
1242    /// Build a trace with a single attempt at `rung`, `total_cost_usd`, and a given logging
1243    /// propensity set on `policy.propensity`. Used to construct a known-distribution sample for
1244    /// IPS correctness assertions.
1245    fn make_propensity_trace(
1246        tenant: &str,
1247        rung: u32,
1248        cost_usd: f64,
1249        propensity: Option<f64>,
1250    ) -> Trace {
1251        let attempt = firstpass_core::Attempt {
1252            rung,
1253            model: "m".to_owned(),
1254            provider: "anthropic".to_owned(),
1255            in_tokens: 10,
1256            out_tokens: 5,
1257            cost_usd,
1258            latency_ms: 5,
1259            gates: vec![],
1260            verdict: Verdict::Pass,
1261        };
1262        let mut trace = Trace {
1263            trace_id: uuid::Uuid::now_v7(),
1264            prev_hash: GENESIS_HASH.to_owned(),
1265            tenant_id: tenant.to_owned(),
1266            session_id: "s".to_owned(),
1267            ts: jiff::Timestamp::now(),
1268            mode: Mode::Enforce,
1269            policy: PolicyRef {
1270                id: "bandit@v1+eps".to_owned(),
1271                explore: rung != 0,
1272                propensity,
1273                mode_profile: None,
1274            },
1275            request: RequestInfo {
1276                api: "anthropic.messages".to_owned(),
1277                prompt_hash: "ph".to_owned(),
1278                features: Features::new(TaskKind::Other),
1279            },
1280            attempts: vec![attempt],
1281            deferred: Vec::new(),
1282            final_: FinalOutcome {
1283                served_rung: Some(rung),
1284                served_from: ServedFrom::Attempt,
1285                total_cost_usd: cost_usd,
1286                gate_cost_usd: 0.0,
1287                total_latency_ms: 5,
1288                escalations: 0,
1289                counterfactual_baseline_usd: cost_usd,
1290                savings_usd: 0.0,
1291            },
1292            probe: None,
1293            rollout: None,
1294            shadow: None,
1295            predicted_pass: None,
1296            elastic: None,
1297        };
1298        trace.recompute_savings();
1299        trace
1300    }
1301
1302    // ── 10. IPS correctness from a known logging policy ───────────────────────
1303    //
1304    // Logging policy: epsilon-greedy, K=2 rungs, greedy=rung 0, epsilon=0.2
1305    //   p(start==0) = (1-0.2)*1 + 0.2/2 = 0.9
1306    //   p(start==1) = (1-0.2)*0 + 0.2/2 = 0.1
1307    //
1308    // Trace set representative of the logging policy:
1309    //   45 traces at rung 0 (cost $0.001, propensity 0.9)
1310    //    5 traces at rung 1 (cost $0.010, propensity 0.1)
1311    //
1312    // Candidate: always start at rung 1 (start_rung = 1).
1313    //   w_i = 1(start==1) / p_i
1314    //   Rung-0 traces: w = 0/0.9 = 0
1315    //   Rung-1 traces: w = 1/0.1 = 10
1316    //
1317    //   sum_wc = 5 * 10 * 0.010 = 0.5
1318    //   IPS  = sum_wc / n = 0.5 / 50 = 0.010  βœ“  (equals true mean cost of rung-1 traces)
1319    //   SNIPS = sum_wc / sum_w = 0.5 / (5*10) = 0.010 βœ“
1320    //   ESS   = (5*10)^2 / (5*10^2) = 2500/500 = 5
1321    #[tokio::test]
1322    async fn ips_correctness_from_known_logging_policy() {
1323        let db = tmp_db();
1324        let (tx, handle) = store::open(&db).unwrap();
1325
1326        for _ in 0..45 {
1327            tx.try_send(make_propensity_trace("t", 0, 0.001, Some(0.9)))
1328                .unwrap();
1329        }
1330        for _ in 0..5 {
1331            tx.try_send(make_propensity_trace("t", 1, 0.010, Some(0.1)))
1332                .unwrap();
1333        }
1334        drop(tx);
1335        handle.await.unwrap();
1336
1337        let report = ips_from_store(&db, "t", 1).unwrap();
1338
1339        assert_eq!(report.n_traces, 50);
1340        assert_eq!(report.n_with_propensity, 50);
1341        // IPS = SNIPS = true mean cost of rung-1-start traces = $0.010.
1342        assert!(
1343            (report.ips_cost - 0.010).abs() < 1e-9,
1344            "IPS cost={} expected 0.010",
1345            report.ips_cost
1346        );
1347        assert!(
1348            (report.snips_cost - 0.010).abs() < 1e-9,
1349            "SNIPS cost={} expected 0.010",
1350            report.snips_cost
1351        );
1352        assert!(
1353            report.ess.is_finite() && report.ess > 0.0,
1354            "ESS must be positive"
1355        );
1356        assert!(
1357            (report.ess - 5.0).abs() < 1e-9,
1358            "ESS={} expected 5.0",
1359            report.ess
1360        );
1361        // No deferred feedback β†’ no correctness signal.
1362        assert_eq!(report.n_correctness_known, 0);
1363        assert!(report.ips_served_failure.is_none());
1364
1365        let _ = std::fs::remove_file(&db);
1366    }
1367
1368    // ── 11. Traces without propensity are excluded from IPS ───────────────────
1369
1370    #[tokio::test]
1371    async fn ips_excludes_traces_without_propensity() {
1372        let db = tmp_db();
1373        let (tx, handle) = store::open(&db).unwrap();
1374
1375        // 5 traces with propensity, 5 without.
1376        for _ in 0..5 {
1377            tx.try_send(make_propensity_trace("t", 0, 0.001, Some(0.9)))
1378                .unwrap();
1379        }
1380        for _ in 0..5 {
1381            tx.try_send(make_propensity_trace("t", 0, 0.001, None))
1382                .unwrap();
1383        }
1384        drop(tx);
1385        handle.await.unwrap();
1386
1387        let report = ips_from_store(&db, "t", 0).unwrap();
1388
1389        assert_eq!(report.n_traces, 10);
1390        assert_eq!(
1391            report.n_with_propensity, 5,
1392            "only traces with propensity count"
1393        );
1394        // All 5 propensity traces start at rung 0 (candidate rung 0): w = 1/0.9 each.
1395        // IPS = sum_wc / n = 5 * (1/0.9) * 0.001 / 5 β‰ˆ 0.001/0.9 β‰ˆ 0.001111
1396        let expected_ips = 0.001_f64 / 0.9;
1397        assert!(
1398            (report.ips_cost - expected_ips).abs() < 1e-9,
1399            "IPS={} expected {expected_ips}",
1400            report.ips_cost
1401        );
1402
1403        let _ = std::fs::remove_file(&db);
1404    }
1405
1406    // ── 12. Empty store / no propensity traces β†’ zero report ─────────────────
1407
1408    #[test]
1409    fn ips_empty_store_returns_zero_report() {
1410        let db = std::path::Path::new("/nonexistent/fp-ips-empty.db");
1411        let report = ips_from_store(db, "t", 0).unwrap();
1412        assert_eq!(report.n_traces, 0);
1413        assert_eq!(report.n_with_propensity, 0);
1414        assert_eq!(report.ips_cost, 0.0);
1415        assert_eq!(report.snips_cost, 0.0);
1416        assert_eq!(report.ess, 0.0);
1417        assert!(report.ips_served_failure.is_none());
1418        // DR zero-trace case.
1419        assert_eq!(report.dr_cost, 0.0);
1420        assert_eq!(report.ci_dr_cost, (0.0, 0.0));
1421    }
1422
1423    // ── DR estimator tests ────────────────────────────────────────────────────
1424
1425    /// Helper: like `make_propensity_trace` but with `TaskKind::CodeEdit` so tests can build
1426    /// a two-context population for the sparse-bucket fallback test.
1427    fn make_propensity_trace_ce(rung: u32, cost: f64, p: Option<f64>) -> Trace {
1428        let mut t = make_propensity_trace("t", rung, cost, p);
1429        t.request.features.task_kind = TaskKind::CodeEdit;
1430        t
1431    }
1432
1433    // ── 13. DR = DM = IPS when all traces are at the candidate rung (p=1.0) ──
1434    //
1435    // Degenerate logging policy: always start at rung 1, propensity = 1.0.
1436    //   wα΅’ = 1/1.0 = 1 for every trace.
1437    //   DM(Other, 1) = mean(costs) = 0.010.
1438    //   DRα΅’ = DM + 1Β·(costα΅’ βˆ’ DM) = costα΅’  β†’  DR = mean(cost) = DM = IPS.
1439    #[tokio::test]
1440    async fn dr_degenerate_propensities_equals_dm_and_ips() {
1441        let db = tmp_db();
1442        let (tx, handle) = store::open(&db).unwrap();
1443
1444        for _ in 0..5 {
1445            tx.try_send(make_propensity_trace("t", 1, 0.010, Some(1.0)))
1446                .unwrap();
1447        }
1448        drop(tx);
1449        handle.await.unwrap();
1450
1451        let r = ips_from_store(&db, "t", 1).unwrap();
1452
1453        // All three estimators agree when the logging policy is degenerate-correct.
1454        assert!(
1455            (r.dr_cost - 0.010).abs() < 1e-9,
1456            "DR={} expected 0.010",
1457            r.dr_cost
1458        );
1459        assert!(
1460            (r.dr_cost - r.ips_cost).abs() < 1e-9,
1461            "DR should equal IPS; DR={} IPS={}",
1462            r.dr_cost,
1463            r.ips_cost
1464        );
1465        assert!(
1466            (r.dr_cost - r.snips_cost).abs() < 1e-9,
1467            "DR should equal SNIPS; DR={} SNIPS={}",
1468            r.dr_cost,
1469            r.snips_cost
1470        );
1471
1472        let _ = std::fs::remove_file(&db);
1473    }
1474
1475    // ── 14. DR correction toward truth under correct propensities ─────────────
1476    //
1477    // Logging policy: Ξ΅-greedy, K=2, greedy=rung 0, Ξ΅=0.2:
1478    //   p(start=0) = 0.9,  p(start=1) = 0.1
1479    //
1480    // Data: 9 traces at rung 0 (cost $0.001, p=0.9), 1 trace at rung 1 (cost $0.010, p=0.1).
1481    //
1482    // Naive mean of ALL logged costs: (9Β·0.001 + 1Β·0.010)/10 = 0.0019  ← biased for
1483    // "always rung 1".  DM(Other, 1) = 0.010 (correct bucket mean from the one observation).
1484    //
1485    // IPS  = (1/0.1)Β·0.010 / 10 = 0.010  βœ“
1486    // DR:
1487    //   9 rung-0 traces: DRα΅’ = DM(Other,1) + 0Β·(…) = 0.010
1488    //   1 rung-1 trace:  DRα΅’ = DM(Other,1) + 10Β·(0.010 βˆ’ 0.010) = 0.010
1489    //   DR = 0.010 βœ“   (not the biased 0.0019)
1490    #[tokio::test]
1491    async fn dr_correction_recovers_true_cost_under_correct_propensities() {
1492        let db = tmp_db();
1493        let (tx, handle) = store::open(&db).unwrap();
1494
1495        for _ in 0..9 {
1496            tx.try_send(make_propensity_trace("t", 0, 0.001, Some(0.9)))
1497                .unwrap();
1498        }
1499        tx.try_send(make_propensity_trace("t", 1, 0.010, Some(0.1)))
1500            .unwrap();
1501        drop(tx);
1502        handle.await.unwrap();
1503
1504        let r = ips_from_store(&db, "t", 1).unwrap();
1505
1506        // Naive logged mean = 0.0019; both DR and IPS correctly recover 0.010.
1507        assert!(
1508            (r.dr_cost - 0.010).abs() < 1e-9,
1509            "DR={} expected 0.010 (not the naive logged mean 0.0019)",
1510            r.dr_cost
1511        );
1512        assert!(
1513            (r.dr_cost - r.ips_cost).abs() < 1e-9,
1514            "DR should equal IPS under correct propensities; DR={} IPS={}",
1515            r.dr_cost,
1516            r.ips_cost
1517        );
1518
1519        let _ = std::fs::remove_file(&db);
1520    }
1521
1522    // ── 15. Sparse-bucket fallback β€” no NaN ───────────────────────────────────
1523    //
1524    // Two task kinds create a cross-context sparse-bucket scenario:
1525    //   β€’ CodeEdit traces are only logged at rung 0  β†’ DM(CE, 1) falls back to rung_mean[1]
1526    //   β€’ Other traces are only logged at rung 1     β†’ DM(Other, 0) falls back to rung_mean[0]
1527    //
1528    // DM fallback values must not NaN; DR must be finite.
1529    //
1530    // Candidate rung = 1.
1531    //   CE  rung-0 (5): w=0; DRα΅’ = DM(CE,1) = rung_mean[1] = 0.020
1532    //   Other rung-1 (5): w=2; DRα΅’ = 0.020 + 2Β·(0.020βˆ’0.020) = 0.020
1533    //   DR = 0.020
1534    #[tokio::test]
1535    async fn dr_sparse_bucket_fallback_no_nan() {
1536        let db = tmp_db();
1537        let (tx, handle) = store::open(&db).unwrap();
1538
1539        for _ in 0..5 {
1540            tx.try_send(make_propensity_trace_ce(0, 0.001, Some(0.5)))
1541                .unwrap();
1542        }
1543        for _ in 0..5 {
1544            tx.try_send(make_propensity_trace("t", 1, 0.020, Some(0.5)))
1545                .unwrap();
1546        }
1547        drop(tx);
1548        handle.await.unwrap();
1549
1550        let r = ips_from_store(&db, "t", 1).unwrap();
1551
1552        assert!(
1553            r.dr_cost.is_finite(),
1554            "DR must be finite even when a context bucket is empty (got NaN/inf)"
1555        );
1556        assert!(
1557            (r.dr_cost - 0.020).abs() < 1e-9,
1558            "DR={} expected 0.020",
1559            r.dr_cost
1560        );
1561
1562        let _ = std::fs::remove_file(&db);
1563    }
1564
1565    // ── 16. DR bootstrap CI is present, finite, and sane ─────────────────────
1566    //
1567    // Verifies that the CI brackets the point estimate and is deterministic.
1568    #[tokio::test]
1569    async fn dr_ci_present_and_sane() {
1570        let db = tmp_db();
1571        let (tx, handle) = store::open(&db).unwrap();
1572
1573        // 20 traces at rung 1, alternating cost 0.001 / 0.002, propensity 0.5.
1574        for i in 0..20u32 {
1575            let cost = if i % 2 == 0 { 0.001 } else { 0.002 };
1576            tx.try_send(make_propensity_trace("t", 1, cost, Some(0.5)))
1577                .unwrap();
1578        }
1579        drop(tx);
1580        handle.await.unwrap();
1581
1582        let r1 = ips_from_store(&db, "t", 1).unwrap();
1583        let r2 = ips_from_store(&db, "t", 1).unwrap();
1584
1585        // Deterministic.
1586        assert_eq!(r1.ci_dr_cost, r2.ci_dr_cost, "DR CI must be deterministic");
1587
1588        // CI must be ordered and finite.
1589        let (lo, hi) = r1.ci_dr_cost;
1590        assert!(lo.is_finite() && hi.is_finite(), "DR CI must be finite");
1591        assert!(lo <= hi, "DR CI must be ordered: lo={lo} hi={hi}");
1592
1593        // CI must bracket the point estimate (within floating-point tolerance).
1594        assert!(
1595            lo <= r1.dr_cost + 1e-9,
1596            "DR CI lo {lo} > dr_cost {}",
1597            r1.dr_cost
1598        );
1599        assert!(
1600            hi >= r1.dr_cost - 1e-9,
1601            "DR CI hi {hi} < dr_cost {}",
1602            r1.dr_cost
1603        );
1604
1605        let _ = std::fs::remove_file(&db);
1606    }
1607}