Skip to main content

samkhya_core/
residual.rs

1//! Residual correction model.
2//!
3//! Optional feedback-driven correction layer. Takes a baseline cardinality
4//! estimate plus a feature vector (query plan + column stats) and returns
5//! a corrected estimate. Trained on observations recorded by
6//! [`crate::feedback`].
7//!
8//! Contracts every backend honors:
9//!
10//! - bounded — output never exceeds the LpBound ceiling ([`crate::lpbound`]); the corrector clamps.
11//! - sub-MB / sub-ms — model footprint and per-estimate latency are the architectural budget.
12//! - optional — engines opt in; with no model attached, samkhya behaves as portable stats + envelope.
13//!
14//! Concrete backends (all behind feature flags):
15//!
16//! - `gbt` — gradient-boosted trees (the `gbt` submodule, gated on the `gbt` cargo feature)
17//! - `additive_gbt` — additive gradient-boosted trees (the `additive` submodule, gated on `additive_gbt`)
18//! - `tabpfn` — foundation-model interface (see the `tabpfn` submodule,
19//!   gated on the `tabpfn_http` cargo feature)
20//! - `llm` — LLM-pluggable corrector backend (see the `llm` submodule,
21//!   gated on the `llm_http` cargo feature). Same wire contract as
22//!   `tabpfn`; the server picks an Anthropic / OpenAI / local-Ollama /
23//!   dummy provider via the `SAMKHYA_LLM_BACKEND` env var.
24//!
25//! # Foundation-model interface (Layer 5)
26//!
27//! The architecture reserves a pluggable backend slot for foundation tabular
28//! models such as TabPFN-2.5 (arXiv 2511.08667). The contract is identical
29//! to every other backend:
30//!
31//! > *feed [`CorrectionFeatures`], receive `Option<u64>` clamped to the
32//! > LpBound ceiling.*
33//!
34//! Two deployment shapes are scaffolded:
35//!
36//! 1. **localhost HTTP** — a Python TabPFN inference server runs out of
37//!    band; samkhya POSTs the feature vector as JSON and reads back an
38//!    `{"estimate": <u64>}` response. Implemented today behind the
39//!    `tabpfn_http` cargo feature (see `tabpfn::TabPfnHttpCorrector`).
40//! 2. **subprocess** — samkhya spawns a Python child, frames JSON over
41//!    stdin/stdout, and keeps the process warm across estimates. Deferred:
42//!    the scaffolding is present (umbrella `tabpfn` feature), the
43//!    transport itself is not implemented in this revision.
44//!
45//! A no-op [`TabPfnStub`] is **always** compiled in, regardless of
46//! features. Its job is purely architectural: downstream code can reference
47//! `TabPfnStub` to mark "TabPFN integration point, currently disabled"
48//! without taking the `tabpfn_http` feature dependency. This reflects the
49//! integration point in every build, so the contract is visible even when
50//! the transport is not.
51//!
52//! Failure policy across all TabPFN backends: any transport error, parse
53//! error, or timeout returns `Ok(None)` (never `Err`). The engine then
54//! falls back cleanly to the native estimate. This is the safety contract;
55//! a remote inference server going down must never surface as a query
56//! failure.
57
58use crate::Result;
59
60/// Emit a single `log::warn!` the first time a plaintext-HTTP URL
61/// pointing at a non-loopback host is configured for an HTTP corrector
62/// backend. See `documents/SECURITY-REVIEW-2026-05-17.md` (H2): such a
63/// URL means features and any embedded baseline estimate travel
64/// unencrypted on the wire. The warning is fire-and-forget — no behaviour
65/// change, so well-configured operators (the typical case, defaults are
66/// localhost) see nothing.
67#[cfg(any(feature = "tabpfn_http", feature = "llm_http"))]
68fn warn_if_remote_plaintext_http(url: &str, backend: &'static str) {
69    let lower = url.to_ascii_lowercase();
70    if !lower.starts_with("http://") {
71        return;
72    }
73    // Pull the host (between "http://" and the next "/" or ":" or end).
74    let rest = &url[7..]; // safe: starts_with confirmed above
75    let host_end = rest.find(['/', ':', '?']).unwrap_or(rest.len());
76    let host = &rest[..host_end];
77    let is_loopback = matches!(host, "127.0.0.1" | "::1" | "localhost")
78        || host.starts_with("[::1]")
79        || host.starts_with("127.");
80    if is_loopback {
81        return;
82    }
83    if std::env::var("SAMKHYA_ALLOW_REMOTE_HTTP").as_deref() == Ok("1") {
84        return;
85    }
86    log::warn!(
87        "samkhya {backend} corrector configured with plaintext HTTP to non-loopback host {host}; \
88         features and baseline_estimate will travel unencrypted. Use https:// or set \
89         SAMKHYA_ALLOW_REMOTE_HTTP=1 to silence this warning."
90    );
91}
92
93/// Feature vector handed to the corrector at estimate time.
94///
95/// Intentionally minimal at v0.0.1: row count + distinct count + null
96/// count + a small set of operator-level features. Will grow as the
97/// feedback-collection surface widens.
98///
99/// # Examples
100///
101/// ```
102/// use samkhya_core::residual::CorrectionFeatures;
103///
104/// let features = CorrectionFeatures {
105///     baseline_estimate: 1000,
106///     left_input_rows: Some(500),
107///     right_input_rows: Some(2000),
108///     predicate_count: 2,
109///     join_depth: 1,
110///     ..Default::default()
111/// };
112/// assert_eq!(features.to_vec().len(), CorrectionFeatures::FEATURE_LEN);
113/// ```
114#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
115pub struct CorrectionFeatures {
116    pub baseline_estimate: u64,
117    pub left_input_rows: Option<u64>,
118    pub right_input_rows: Option<u64>,
119    pub left_distinct: Option<u64>,
120    pub right_distinct: Option<u64>,
121    pub predicate_count: u32,
122    pub join_depth: u32,
123}
124
125impl CorrectionFeatures {
126    /// Flatten the feature struct into a fixed-length numeric vector for a
127    /// regression model. `Option<u64>` slots are zero-filled when absent —
128    /// callers should treat zero as "unknown" rather than "literally zero
129    /// rows", which is the convention the corrector is trained against.
130    ///
131    /// Layout (stable; new features must be appended, never reordered):
132    ///
133    /// 0. `baseline_estimate`
134    /// 1. `left_input_rows`  (0 if `None`)
135    /// 2. `right_input_rows` (0 if `None`)
136    /// 3. `left_distinct`    (0 if `None`)
137    /// 4. `right_distinct`   (0 if `None`)
138    /// 5. `predicate_count`
139    /// 6. `join_depth`
140    ///
141    /// # Examples
142    ///
143    /// ```
144    /// use samkhya_core::residual::CorrectionFeatures;
145    ///
146    /// let f = CorrectionFeatures {
147    ///     baseline_estimate: 100,
148    ///     left_input_rows: Some(10),
149    ///     predicate_count: 3,
150    ///     ..Default::default()
151    /// };
152    /// let v = f.to_vec();
153    /// assert_eq!(v[0], 100.0);
154    /// assert_eq!(v[1], 10.0);
155    /// assert_eq!(v[2], 0.0); // None → 0
156    /// assert_eq!(v[5], 3.0);
157    /// ```
158    pub fn to_vec(&self) -> Vec<f64> {
159        vec![
160            self.baseline_estimate as f64,
161            self.left_input_rows.unwrap_or(0) as f64,
162            self.right_input_rows.unwrap_or(0) as f64,
163            self.left_distinct.unwrap_or(0) as f64,
164            self.right_distinct.unwrap_or(0) as f64,
165            f64::from(self.predicate_count),
166            f64::from(self.join_depth),
167        ]
168    }
169
170    /// Number of entries [`to_vec`](Self::to_vec) produces.
171    pub const FEATURE_LEN: usize = 7;
172}
173
174/// A pluggable corrector. Engines call [`Corrector::correct`] on every
175/// estimate that passes through samkhya's optimizer hook.
176///
177/// Returning `Ok(None)` lets the engine fall back to the baseline estimate;
178/// returning `Ok(Some(_))` overrides it (subject to the LpBound envelope).
179///
180/// # Examples
181///
182/// ```
183/// use samkhya_core::residual::{CorrectionFeatures, Corrector, IdentityCorrector};
184///
185/// let corrector = IdentityCorrector;
186/// let features = CorrectionFeatures {
187///     baseline_estimate: 42,
188///     ..Default::default()
189/// };
190/// // The identity corrector passes the baseline through unchanged.
191/// assert_eq!(corrector.correct(&features).unwrap(), Some(42));
192/// ```
193pub trait Corrector: Send + Sync {
194    /// Return a corrected estimate, or `None` to fall back to the baseline.
195    fn correct(&self, features: &CorrectionFeatures) -> Result<Option<u64>>;
196
197    /// Stable identifier for logging / model-version tracking.
198    fn name(&self) -> &'static str;
199}
200
201/// Default zero-cost corrector: passes the baseline through unchanged.
202///
203/// Used when no feedback history exists yet (cold start) or when the
204/// caller opts out of feedback-driven correction entirely.
205///
206/// # Examples
207///
208/// ```
209/// use samkhya_core::residual::{CorrectionFeatures, Corrector, IdentityCorrector};
210///
211/// let c = IdentityCorrector;
212/// let f = CorrectionFeatures { baseline_estimate: 1234, ..Default::default() };
213/// assert_eq!(c.correct(&f).unwrap(), Some(1234));
214/// assert_eq!(c.name(), "identity");
215/// ```
216pub struct IdentityCorrector;
217
218impl Corrector for IdentityCorrector {
219    fn correct(&self, features: &CorrectionFeatures) -> Result<Option<u64>> {
220        Ok(Some(features.baseline_estimate))
221    }
222
223    fn name(&self) -> &'static str {
224        "identity"
225    }
226}
227
228/// No-op stub for the foundation-model interface (Layer 5).
229///
230/// Always compiled, regardless of cargo features. Returns `Ok(None)` from
231/// every call, signalling the engine to fall back to its native estimate.
232///
233/// The point of an always-on stub is architectural: it lets downstream
234/// callers reference `TabPfnStub` to mark "TabPFN integration point,
235/// currently disabled" without taking the `tabpfn_http` feature
236/// dependency. The integration shape is visible in every build.
237///
238/// To wire in a real foundation-model backend, swap this for
239/// `tabpfn::TabPfnHttpCorrector` (gated on `tabpfn_http`) or a future
240/// subprocess adapter. The trait contract is identical, so the swap is
241/// a one-line change at the call site.
242///
243/// # Examples
244///
245/// ```
246/// use samkhya_core::residual::{CorrectionFeatures, Corrector, TabPfnStub};
247///
248/// let stub = TabPfnStub;
249/// // Stub always returns Ok(None) — the engine falls back to its native estimate.
250/// let f = CorrectionFeatures { baseline_estimate: 999, ..Default::default() };
251/// assert_eq!(stub.correct(&f).unwrap(), None);
252/// assert_eq!(stub.name(), "tabpfn-stub");
253/// ```
254pub struct TabPfnStub;
255
256impl Corrector for TabPfnStub {
257    fn correct(&self, _features: &CorrectionFeatures) -> Result<Option<u64>> {
258        // Deliberately `None`: the integration point is wired but
259        // disabled. The engine falls back to the native estimate.
260        Ok(None)
261    }
262
263    fn name(&self) -> &'static str {
264        "tabpfn-stub"
265    }
266}
267
268#[cfg(feature = "gbt")]
269pub mod gbt {
270    //! Gradient-boosted-tree residual corrector.
271    //!
272    //! Wraps the `gbdt` crate (Baidu / mesalock-linux,
273    //! <https://github.com/mesalock-linux/gbdt-rs>) — pure-Rust, no native
274    //! deps, builds on stable Rust 1.94 / edition 2024. Compiled in only
275    //! when the `gbt` cargo feature is enabled.
276    //!
277    //! Training target is `log(actual_rows / est_rows)` — the
278    //! multiplicative correction ratio in log-space. At prediction time
279    //! we exponentiate and multiply through the baseline, then clamp to
280    //! the configured LpBound ceiling via
281    //! [`crate::lpbound::saturating_clamp`] so the corrector cannot ever
282    //! violate the envelope contract.
283    //!
284    //! Observations with `est_rows == 0` or `actual_rows == 0` are
285    //! silently dropped (log of zero is undefined); we do not invent a
286    //! Laplace-style smoothing constant at the corrector layer.
287
288    use gbdt::config::{Config, Loss};
289    use gbdt::decision_tree::{Data, DataVec};
290    use gbdt::gradient_boost::GBDT;
291
292    use super::{CorrectionFeatures, Corrector};
293    use crate::feedback::{Observation, PlanObservation};
294    use crate::lpbound::saturating_clamp;
295    use crate::{Error, Result};
296
297    /// Tunables for [`GbtCorrector::train`]. Defaults are an MVP starting
298    /// point: shallow trees, modest depth, square-error loss.
299    #[derive(Debug, Clone)]
300    pub struct GbtOptions {
301        /// Shrinkage / learning rate applied to each tree's contribution.
302        pub learning_rate: f64,
303        /// Max depth of each regression tree. Root is depth 0.
304        pub max_depth: u32,
305        /// Number of boosting iterations (one tree per iteration).
306        pub num_trees: u32,
307        /// Inclusive upper bound applied to every corrected estimate.
308        /// Use `u64::MAX` to disable (the trait signature has no ceiling
309        /// slot, so we store it here at train time).
310        pub ceiling: u64,
311        /// Minimum samples per leaf — guards against overfitting tiny
312        /// feedback histories.
313        pub min_leaf_size: usize,
314    }
315
316    impl Default for GbtOptions {
317        fn default() -> Self {
318            Self {
319                learning_rate: 0.1,
320                max_depth: 4,
321                num_trees: 50,
322                ceiling: u64::MAX,
323                min_leaf_size: 1,
324            }
325        }
326    }
327
328    /// Trained GBT-backed residual corrector.
329    pub struct GbtCorrector {
330        model: GBDT,
331        ceiling: u64,
332        training_rows: usize,
333    }
334
335    impl GbtCorrector {
336        /// Train a corrector from a slice of [`Observation`]s.
337        ///
338        /// Returns [`Error::Feedback`] if the observation slice is empty,
339        /// or if every observation is unusable (zero est_rows or zero
340        /// actual_rows). Non-positive-ratio observations are silently
341        /// filtered, matching the convention in [`Observation::q_error`].
342        pub fn train(observations: &[Observation], options: GbtOptions) -> Result<Self> {
343            if observations.is_empty() {
344                return Err(Error::Feedback(
345                    "cannot train GbtCorrector: observation slice is empty".into(),
346                ));
347            }
348
349            let mut training: DataVec = Vec::with_capacity(observations.len());
350            for obs in observations {
351                if obs.est_rows == 0 || obs.actual_rows == 0 {
352                    continue;
353                }
354                // Reconstruct a feature vector from the observation. The
355                // feedback table doesn't yet carry full plan features, so
356                // we synthesize the minimal `baseline_estimate`-only
357                // vector. As `Observation` gains columns the mapping
358                // below should grow in lockstep with `CorrectionFeatures`.
359                let features = CorrectionFeatures {
360                    baseline_estimate: obs.est_rows,
361                    ..Default::default()
362                };
363                let feature_f32: Vec<f32> =
364                    features.to_vec().into_iter().map(|v| v as f32).collect();
365                let ratio_log = (obs.actual_rows as f64 / obs.est_rows as f64).ln() as f32;
366                training.push(Data::new_training_data(feature_f32, 1.0, ratio_log, None));
367            }
368
369            if training.is_empty() {
370                return Err(Error::Feedback(
371                    "cannot train GbtCorrector: all observations had zero est or actual rows"
372                        .into(),
373                ));
374            }
375
376            let mut cfg = Config::new();
377            cfg.set_feature_size(CorrectionFeatures::FEATURE_LEN);
378            cfg.set_max_depth(options.max_depth);
379            cfg.set_iterations(options.num_trees as usize);
380            cfg.set_shrinkage(options.learning_rate as f32);
381            cfg.set_min_leaf_size(options.min_leaf_size);
382            cfg.set_loss(&loss_name(Loss::SquaredError));
383
384            let rows = training.len();
385            let mut model = GBDT::new(&cfg);
386            model.fit(&mut training);
387
388            Ok(Self {
389                model,
390                ceiling: options.ceiling,
391                training_rows: rows,
392            })
393        }
394
395        /// Train from observations that carry the plan features the
396        /// corrector will actually see at inference time.
397        ///
398        /// Prefer this over [`train`](Self::train). That method has to
399        /// synthesise a feature vector with only `baseline_estimate`
400        /// populated, because [`Observation`] carries nothing else — so the
401        /// remaining six features are constant across the whole training
402        /// set, no tree ever splits on them, and the model is effectively
403        /// one-dimensional while the adapter feeds it seven live features.
404        /// The mismatch is silent: nothing fails, the corrector is just
405        /// blind to everything except the baseline.
406        ///
407        /// Observations with `baseline_estimate == 0` or `actual_rows == 0`
408        /// are dropped, since the target is `log(actual / baseline)`.
409        ///
410        /// # Examples
411        ///
412        /// ```
413        /// use samkhya_core::feedback::PlanObservation;
414        /// use samkhya_core::residual::CorrectionFeatures;
415        /// use samkhya_core::residual::gbt::{GbtCorrector, GbtOptions};
416        ///
417        /// let observations: Vec<PlanObservation> = (1..40u64)
418        ///     .map(|i| PlanObservation {
419        ///         template_hash: "t".into(),
420        ///         plan_fingerprint: "p".into(),
421        ///         features: CorrectionFeatures {
422        ///             baseline_estimate: i * 10,
423        ///             join_depth: (i % 4) as u32,
424        ///             ..Default::default()
425        ///         },
426        ///         actual_rows: i * 25,
427        ///         latency_ms: None,
428        ///     })
429        ///     .collect();
430        ///
431        /// let corrector = GbtCorrector::train_on_plans(&observations, GbtOptions::default())
432        ///     .expect("trains");
433        /// assert_eq!(corrector.training_rows(), 39);
434        /// ```
435        pub fn train_on_plans(
436            observations: &[PlanObservation],
437            options: GbtOptions,
438        ) -> Result<Self> {
439            if observations.is_empty() {
440                return Err(Error::Feedback(
441                    "cannot train GbtCorrector: observation slice is empty".into(),
442                ));
443            }
444
445            let mut training: DataVec = Vec::with_capacity(observations.len());
446            for obs in observations {
447                if obs.features.baseline_estimate == 0 || obs.actual_rows == 0 {
448                    continue;
449                }
450                let feature_f32: Vec<f32> = obs
451                    .features
452                    .to_vec()
453                    .into_iter()
454                    .map(|v| v as f32)
455                    .collect();
456                let ratio_log =
457                    (obs.actual_rows as f64 / obs.features.baseline_estimate as f64).ln() as f32;
458                training.push(Data::new_training_data(feature_f32, 1.0, ratio_log, None));
459            }
460
461            if training.is_empty() {
462                return Err(Error::Feedback(
463                    "cannot train GbtCorrector: every observation had a zero baseline or actual"
464                        .into(),
465                ));
466            }
467
468            let rows = training.len();
469            let mut cfg = Config::new();
470            cfg.set_feature_size(CorrectionFeatures::FEATURE_LEN);
471            cfg.set_max_depth(options.max_depth);
472            cfg.set_iterations(options.num_trees as usize);
473            cfg.set_shrinkage(options.learning_rate as f32);
474            cfg.set_min_leaf_size(options.min_leaf_size);
475            cfg.set_loss(&loss_name(Loss::SquaredError));
476
477            let mut model = GBDT::new(&cfg);
478            model.fit(&mut training);
479
480            Ok(Self {
481                model,
482                ceiling: options.ceiling,
483                training_rows: rows,
484            })
485        }
486
487        /// Persist the trained model so a later process can evaluate with
488        /// it instead of retraining.
489        ///
490        /// Keeping training and evaluation in separate processes is what
491        /// makes an honest held-out measurement possible: the model cannot
492        /// see the evaluation queries if it was frozen before they ran.
493        pub fn save(&self, path: impl AsRef<std::path::Path>) -> Result<()> {
494            let path = path.as_ref();
495            let name = path.to_str().ok_or_else(|| {
496                Error::Feedback(format!("model path is not valid UTF-8: {}", path.display()))
497            })?;
498            self.model
499                .save_model(name)
500                .map_err(|e| Error::Feedback(format!("could not save GBT model: {e}")))
501        }
502
503        /// Load a model persisted by [`save`](Self::save), applying
504        /// `ceiling` as the clamp.
505        ///
506        /// The ceiling is supplied at load time rather than stored with the
507        /// model because it is a property of the query being bounded, not
508        /// of the model.
509        pub fn load(path: impl AsRef<std::path::Path>, ceiling: u64) -> Result<Self> {
510            let path = path.as_ref();
511            let name = path.to_str().ok_or_else(|| {
512                Error::Feedback(format!("model path is not valid UTF-8: {}", path.display()))
513            })?;
514            let model = GBDT::load_model(name)
515                .map_err(|e| Error::Feedback(format!("could not load GBT model: {e}")))?;
516            Ok(Self {
517                model,
518                ceiling,
519                training_rows: 0,
520            })
521        }
522
523        /// Number of observations the model was fitted on. Zero for a model
524        /// restored by [`load`](Self::load), which does not carry it.
525        pub fn training_rows(&self) -> usize {
526            self.training_rows
527        }
528
529        /// Predict the log-ratio correction for a single feature vector.
530        /// Exposed for diagnostics / unit tests; the production path is
531        /// [`Corrector::correct`].
532        pub fn predict_log_ratio(&self, features: &CorrectionFeatures) -> f64 {
533            let feature_f32: Vec<f32> = features.to_vec().into_iter().map(|v| v as f32).collect();
534            let probe: DataVec = vec![Data::new_test_data(feature_f32, None)];
535            let preds = self.model.predict(&probe);
536            preds.first().copied().unwrap_or(0.0) as f64
537        }
538
539        /// Configured upper bound. Set at training time; the trait method
540        /// [`Corrector::correct`] enforces it via `saturating_clamp`.
541        pub fn ceiling(&self) -> u64 {
542            self.ceiling
543        }
544    }
545
546    impl Corrector for GbtCorrector {
547        fn correct(&self, features: &CorrectionFeatures) -> Result<Option<u64>> {
548            let log_ratio = self.predict_log_ratio(features);
549            let ratio = log_ratio.exp();
550            let scaled = features.baseline_estimate as f64 * ratio;
551            Ok(Some(saturating_clamp(scaled, self.ceiling)))
552        }
553
554        fn name(&self) -> &'static str {
555            "gbt"
556        }
557    }
558
559    /// `gbdt::config::Config::set_loss` takes a string; this is the
560    /// canonical spelling for square-error in that crate.
561    fn loss_name(loss: Loss) -> String {
562        gbdt::config::loss2string(&loss)
563    }
564}
565
566#[cfg(feature = "additive_gbt")]
567pub mod additive {
568    //! Additive gradient-boosted-tree residual corrector.
569    //!
570    //! Sibling backend to [`super::gbt`]. The multiplicative form trains on
571    //! `log(actual / baseline_estimate)` and applies the correction as
572    //! `baseline * exp(predicted)`. That model is structurally trapped at
573    //! zero whenever the engine hands us `baseline_estimate = 0` — the
574    //! q=∞ regime where the upstream estimator has completely collapsed
575    //! (a common DataFusion 46 symptom on chained joins).
576    //!
577    //! The additive backend sidesteps that trap by training the model to
578    //! predict the **absolute** `actual_rows` from the full
579    //! [`CorrectionFeatures`] vector (all 7 features, not just the
580    //! baseline). The prediction is clamped to a non-negative integer and
581    //! then to the configured LpBound ceiling via
582    //! [`crate::lpbound::saturating_clamp`], so the envelope contract is
583    //! preserved.
584    //!
585    //! Cargo feature: `additive_gbt`. Independent of the `gbt` feature —
586    //! they can be enabled separately or together.
587
588    use gbdt::config::{Config, Loss};
589    use gbdt::decision_tree::{Data, DataVec};
590    use gbdt::gradient_boost::GBDT;
591    use std::sync::Mutex;
592
593    use super::{CorrectionFeatures, Corrector};
594    use crate::feedback::Observation;
595    use crate::lpbound::saturating_clamp;
596    use crate::{Error, Result};
597
598    /// Tunables for [`AdditiveGbtCorrector::train`]. Defaults mirror
599    /// [`super::gbt::GbtOptions`] so the two backends are
600    /// drop-in-comparable when benchmarking.
601    #[derive(Debug, Clone)]
602    pub struct AdditiveGbtOptions {
603        /// Shrinkage / learning rate applied to each tree's contribution.
604        pub learning_rate: f64,
605        /// Max depth of each regression tree. Root is depth 0.
606        pub max_depth: u32,
607        /// Number of boosting iterations (one tree per iteration).
608        pub num_trees: u32,
609        /// Inclusive upper bound applied to every corrected estimate.
610        /// Use `u64::MAX` to disable.
611        pub ceiling: u64,
612        /// Minimum samples per leaf — guards against overfitting tiny
613        /// feedback histories.
614        pub min_leaf_size: usize,
615    }
616
617    impl Default for AdditiveGbtOptions {
618        fn default() -> Self {
619            Self {
620                learning_rate: 0.1,
621                max_depth: 4,
622                num_trees: 50,
623                ceiling: u64::MAX,
624                min_leaf_size: 1,
625            }
626        }
627    }
628
629    /// Trained additive GBT corrector. Predicts absolute row counts.
630    ///
631    /// The model is wrapped in a [`Mutex`] because `gbdt::GBDT::predict`
632    /// takes `&mut self` on some configurations; the lock is held only
633    /// for the prediction call and is uncontended in the common single-
634    /// threaded estimate path.
635    pub struct AdditiveGbtCorrector {
636        model: Mutex<GBDT>,
637        ceiling: u64,
638    }
639
640    impl AdditiveGbtCorrector {
641        /// Train an additive corrector from a slice of [`Observation`]s.
642        ///
643        /// Returns [`Error::Feedback`] if the observation slice is empty.
644        /// Unlike the multiplicative backend, observations with
645        /// `est_rows == 0` are **kept** — they are precisely the q=∞
646        /// regime this backend exists to handle. Observations with
647        /// `actual_rows == 0` are also kept (a true-zero output is a
648        /// valid signal for an additive model).
649        pub fn train(observations: &[Observation], options: AdditiveGbtOptions) -> Result<Self> {
650            if observations.is_empty() {
651                return Err(Error::Feedback(
652                    "cannot train AdditiveGbtCorrector: observation slice is empty".into(),
653                ));
654            }
655
656            let mut training: DataVec = Vec::with_capacity(observations.len());
657            for obs in observations {
658                // Reconstruct a feature vector from the observation. The
659                // feedback table doesn't yet carry the full plan-shape
660                // feature set, so we synthesize from `est_rows`. As
661                // `Observation` gains columns, mirror the additions here.
662                let features = CorrectionFeatures {
663                    baseline_estimate: obs.est_rows,
664                    ..Default::default()
665                };
666                let feature_f32: Vec<f32> =
667                    features.to_vec().into_iter().map(|v| v as f32).collect();
668                let target = obs.actual_rows as f32;
669                training.push(Data::new_training_data(feature_f32, 1.0, target, None));
670            }
671
672            // Empty observations are caught above; the synthesized
673            // training set here is always non-empty.
674            debug_assert!(!training.is_empty());
675
676            let mut cfg = Config::new();
677            cfg.set_feature_size(CorrectionFeatures::FEATURE_LEN);
678            cfg.set_max_depth(options.max_depth);
679            cfg.set_iterations(options.num_trees as usize);
680            cfg.set_shrinkage(options.learning_rate as f32);
681            cfg.set_min_leaf_size(options.min_leaf_size);
682            cfg.set_loss(&gbdt::config::loss2string(&Loss::SquaredError));
683
684            let mut model = GBDT::new(&cfg);
685            model.fit(&mut training);
686
687            Ok(Self {
688                model: Mutex::new(model),
689                ceiling: options.ceiling,
690            })
691        }
692
693        /// Predict the absolute row count for a feature vector.
694        /// Exposed for diagnostics; the production path is
695        /// [`Corrector::correct`].
696        pub fn predict_rows(&self, features: &CorrectionFeatures) -> f64 {
697            let feature_f32: Vec<f32> = features.to_vec().into_iter().map(|v| v as f32).collect();
698            let probe: DataVec = vec![Data::new_test_data(feature_f32, None)];
699            let model = self.model.lock().expect("AdditiveGbtCorrector model lock");
700            let preds = model.predict(&probe);
701            preds.first().copied().unwrap_or(0.0) as f64
702        }
703
704        /// Configured upper bound. Set at training time; the trait method
705        /// [`Corrector::correct`] enforces it via `saturating_clamp`.
706        pub fn ceiling(&self) -> u64 {
707            self.ceiling
708        }
709    }
710
711    impl Corrector for AdditiveGbtCorrector {
712        fn correct(&self, features: &CorrectionFeatures) -> Result<Option<u64>> {
713            let raw = self.predict_rows(features).max(0.0);
714            Ok(Some(saturating_clamp(raw, self.ceiling)))
715        }
716
717        fn name(&self) -> &'static str {
718            "additive_gbt"
719        }
720    }
721}
722
723#[cfg(feature = "tabpfn_http")]
724pub mod tabpfn {
725    //! Foundation-model interface — HTTP transport.
726    //!
727    //! Posts a [`super::CorrectionFeatures`] vector as JSON to a
728    //! user-configured endpoint (e.g., a Python TabPFN inference server
729    //! listening on `http://localhost:8765/infer`), parses an
730    //! `{"estimate": <u64>}` reply, and clamps the result to the LpBound
731    //! ceiling via [`crate::lpbound::saturating_clamp`].
732    //!
733    //! Transport: pure-Rust `ureq` (rustls-only, no OpenSSL). Compiled in
734    //! only when the `tabpfn_http` cargo feature is enabled.
735    //!
736    //! # Safety contract
737    //!
738    //! Any failure — DNS, connection refused, HTTP non-2xx, body parse
739    //! error, timeout — returns `Ok(None)`. The engine falls back to the
740    //! native estimate. We never propagate transport errors to the
741    //! optimizer hot path; a remote inference server going down must not
742    //! surface as a query failure.
743    //!
744    //! Note on naming: this is *the foundation-model interface*, not a
745    //! "learned" or "AI" feature. The corrector is a pluggable backend
746    //! behind the same `Corrector` trait as every other backend in this
747    //! module.
748    //!
749    //! # Wire format
750    //!
751    //! Request body (JSON):
752    //!
753    //! ```json
754    //! {
755    //!   "features": [<f64>, <f64>, ...],
756    //!   "baseline_estimate": <u64>
757    //! }
758    //! ```
759    //!
760    //! Response body (JSON):
761    //!
762    //! ```json
763    //! { "estimate": <u64> }
764    //! ```
765    //!
766    //! Any extra fields in the response are ignored, so server
767    //! implementations are free to add diagnostics without breaking the
768    //! client.
769    //!
770    //! # See also
771    //!
772    //! - [`super::TabPfnStub`] — always-on no-op for the same integration
773    //!   slot, no transport dependency.
774
775    use serde::{Deserialize, Serialize};
776    use std::time::Duration;
777
778    use super::{CorrectionFeatures, Corrector};
779    use crate::Result;
780    use crate::lpbound::saturating_clamp;
781
782    /// Configuration for [`TabPfnHttpCorrector`].
783    #[derive(Debug, Clone)]
784    pub struct TabPfnHttpOptions {
785        /// Inference endpoint URL. The corrector POSTs here on every
786        /// `correct()` call. Example: `http://localhost:8765/infer`.
787        pub base_url: String,
788        /// Per-request timeout. Applies independently to the connect and
789        /// read phases. Bounded by the architecture's sub-ms budget for
790        /// the production path, but configurable so users can dial it up
791        /// for diagnostics.
792        pub timeout_ms: u64,
793        /// Inclusive upper bound applied to every corrected estimate via
794        /// [`saturating_clamp`]. The Layer 3 safety guarantee — corrections
795        /// can never exceed this regardless of what the remote backend
796        /// returns. Use `u64::MAX` to disable.
797        pub ceiling: u64,
798    }
799
800    impl Default for TabPfnHttpOptions {
801        fn default() -> Self {
802            Self {
803                base_url: "http://localhost:8765/infer".into(),
804                timeout_ms: 50,
805                ceiling: u64::MAX,
806            }
807        }
808    }
809
810    /// JSON request body sent to the inference endpoint.
811    #[derive(Serialize)]
812    struct InferRequest<'a> {
813        features: &'a [f64],
814        baseline_estimate: u64,
815    }
816
817    /// JSON response body. Extra fields are ignored.
818    #[derive(Deserialize)]
819    struct InferResponse {
820        estimate: u64,
821    }
822
823    /// HTTP-backed foundation-model corrector.
824    ///
825    /// Holds a tiny client config and a base URL. The `ureq` agent is
826    /// constructed per-call: the per-estimate cost is dominated by network
827    /// round-trip, not agent allocation, and per-call agents keep the
828    /// struct cheaply `Send + Sync` without interior mutability.
829    pub struct TabPfnHttpCorrector {
830        options: TabPfnHttpOptions,
831    }
832
833    impl TabPfnHttpCorrector {
834        /// Build a corrector from explicit options.
835        pub fn new(options: TabPfnHttpOptions) -> Self {
836            super::warn_if_remote_plaintext_http(&options.base_url, "tabpfn_http");
837            Self { options }
838        }
839
840        /// Convenience constructor: default options with the supplied URL.
841        pub fn with_url(base_url: impl Into<String>) -> Self {
842            let opts = TabPfnHttpOptions {
843                base_url: base_url.into(),
844                ..TabPfnHttpOptions::default()
845            };
846            super::warn_if_remote_plaintext_http(&opts.base_url, "tabpfn_http");
847            Self { options: opts }
848        }
849
850        /// Configured options (for diagnostics / logging).
851        pub fn options(&self) -> &TabPfnHttpOptions {
852            &self.options
853        }
854
855        /// Attempt one inference call. Returns `None` on any failure
856        /// (network, parse, non-2xx). The `correct()` trait method wraps
857        /// this and applies the LpBound clamp.
858        fn try_infer(&self, features: &CorrectionFeatures) -> Option<u64> {
859            let feature_vec = features.to_vec();
860            let payload = InferRequest {
861                features: &feature_vec,
862                baseline_estimate: features.baseline_estimate,
863            };
864
865            let timeout = Duration::from_millis(self.options.timeout_ms);
866            let agent = ureq::AgentBuilder::new()
867                .timeout_connect(timeout)
868                .timeout_read(timeout)
869                .timeout_write(timeout)
870                .build();
871
872            let response = match agent.post(&self.options.base_url).send_json(&payload) {
873                Ok(r) => r,
874                Err(err) => {
875                    // Map every transport error to None and log at debug.
876                    // The Error::Feedback diagnostic carries the URL plus
877                    // the underlying message so callers tailing logs can
878                    // see what failed without us aborting the query.
879                    log::debug!(
880                        "tabpfn_http: request to {} failed: {}",
881                        self.options.base_url,
882                        err
883                    );
884                    return None;
885                }
886            };
887
888            match response.into_json::<InferResponse>() {
889                Ok(body) => Some(body.estimate),
890                Err(err) => {
891                    log::debug!(
892                        "tabpfn_http: response from {} failed to parse: {}",
893                        self.options.base_url,
894                        err
895                    );
896                    None
897                }
898            }
899        }
900    }
901
902    impl Corrector for TabPfnHttpCorrector {
903        fn correct(&self, features: &CorrectionFeatures) -> Result<Option<u64>> {
904            // Safety contract: every failure returns Ok(None), not Err.
905            // The engine then transparently falls back to the native
906            // estimate. We use Result here to honour the trait shape and
907            // to keep a door open for future *non-fallback* error modes
908            // (e.g. a deliberate misconfiguration check), but on the hot
909            // path failures are absorbed.
910            let Some(raw) = self.try_infer(features) else {
911                return Ok(None);
912            };
913            Ok(Some(saturating_clamp(raw as f64, self.options.ceiling)))
914        }
915
916        fn name(&self) -> &'static str {
917            "tabpfn-http"
918        }
919    }
920}
921
922#[cfg(feature = "llm_http")]
923pub mod llm {
924    //! LLM-pluggable corrector backend — HTTP transport.
925    //!
926    //! Posts a [`super::CorrectionFeatures`] vector as JSON to a
927    //! user-configured endpoint (e.g., a Python LLM inference server
928    //! listening on `http://localhost:8766/infer`) and parses an
929    //! `{"estimate": <u64>}` reply. The server-side LLM provider
930    //! (Anthropic, OpenAI, local Ollama, dummy) is selected by the
931    //! `SAMKHYA_LLM_BACKEND` env var on the server process — the wire
932    //! contract is identical regardless of which provider is configured.
933    //!
934    //! Transport: pure-Rust `ureq` (rustls-only, no OpenSSL). Compiled in
935    //! only when the `llm_http` cargo feature is enabled.
936    //!
937    //! # Naming
938    //!
939    //! This is *the LLM-pluggable corrector backend* — a transport-level
940    //! integration that lets a foundation language model serve as the
941    //! cardinality corrector behind the same `Corrector` trait as every
942    //! other backend in this module. It is **not** an "AI", "adaptive",
943    //! or "learned" feature; the samkhya envelope still dominates the
944    //! safety contract and the LLM is strictly an opt-in pluggable
945    //! backend. The default samkhya build does not pull this in.
946    //!
947    //! # Safety contract
948    //!
949    //! Any failure — DNS, connection refused, HTTP non-2xx, body parse
950    //! error, timeout — returns `Ok(None)`. The engine falls back to the
951    //! native estimate. We never propagate transport errors to the
952    //! optimizer hot path; a remote inference server going down must not
953    //! surface as a query failure. Mirrors the
954    //! [`super::tabpfn::TabPfnHttpCorrector`] contract exactly.
955    //!
956    //! # Wire format
957    //!
958    //! Request body (JSON):
959    //!
960    //! ```json
961    //! {
962    //!   "features": [<f64>, <f64>, ...],
963    //!   "baseline_estimate": <u64>
964    //! }
965    //! ```
966    //!
967    //! Response body (JSON):
968    //!
969    //! ```json
970    //! { "estimate": <u64> }
971    //! ```
972    //!
973    //! Any extra fields in the response are ignored, so server
974    //! implementations are free to add diagnostics (e.g., the LLM's raw
975    //! text reply, parse-status flags) without breaking the client.
976    //!
977    //! # Latency expectations
978    //!
979    //! LLM round-trips are 2–3 orders of magnitude slower than the TabPFN
980    //! tier (P95 in the 0.3–2 s range vs. ~30 ms for TabPFN). The default
981    //! per-request timeout is therefore 2 000 ms (vs. 50 ms for TabPFN),
982    //! with a 60 s hard cap available for cold-cache diagnostics. The
983    //! `llm_http` backend is intended for *offline / overnight*
984    //! re-validation and schema-introspection use cases, not the online
985    //! query hot path. See `bench-results/19_llm_corrector.md` §6 for
986    //! routing guidance.
987
988    use serde::{Deserialize, Serialize};
989    use std::time::Duration;
990
991    use super::{CorrectionFeatures, Corrector};
992    use crate::Result;
993    use crate::lpbound::saturating_clamp;
994
995    /// Default per-request timeout for the LLM HTTP backend (milliseconds).
996    /// LLMs are 2–3 orders of magnitude slower than TabPFN; the 2 s
997    /// default is the smallest budget that consistently covers warm-cache
998    /// Anthropic Claude / OpenAI GPT-4o-mini calls without spurious
999    /// timeouts in measurement.
1000    pub const DEFAULT_TIMEOUT_MS: u64 = 2_000;
1001
1002    /// Hard per-request ceiling (milliseconds). Constructors that accept
1003    /// a `timeout_ms` saturate to this value so a misconfigured caller
1004    /// cannot pin the optimizer for longer than 60 s on a single call.
1005    pub const MAX_TIMEOUT_MS: u64 = 60_000;
1006
1007    /// Default inference endpoint. Distinct from the TabPFN default port
1008    /// (`8765`) so an operator can run both servers side-by-side without
1009    /// collision.
1010    pub const DEFAULT_URL: &str = "http://127.0.0.1:8766/infer";
1011
1012    /// Configuration for [`LlmHttpCorrector`].
1013    #[derive(Debug, Clone)]
1014    pub struct LlmHttpOptions {
1015        /// Inference endpoint URL. The corrector POSTs here on every
1016        /// `correct()` call. Example: `http://localhost:8766/infer`.
1017        pub base_url: String,
1018        /// Per-request timeout. Applies to connect, read, and write
1019        /// phases. Capped at [`MAX_TIMEOUT_MS`] so a misconfigured caller
1020        /// cannot stall the optimizer indefinitely.
1021        pub timeout_ms: u64,
1022        /// Inclusive upper bound applied to every corrected estimate via
1023        /// [`saturating_clamp`]. The Layer 3 safety guarantee —
1024        /// corrections can never exceed this regardless of what the
1025        /// remote LLM returns. Use `u64::MAX` to disable.
1026        pub ceiling: u64,
1027    }
1028
1029    impl Default for LlmHttpOptions {
1030        fn default() -> Self {
1031            Self {
1032                base_url: DEFAULT_URL.into(),
1033                timeout_ms: DEFAULT_TIMEOUT_MS,
1034                ceiling: u64::MAX,
1035            }
1036        }
1037    }
1038
1039    /// JSON request body sent to the inference endpoint.
1040    #[derive(Serialize)]
1041    struct InferRequest<'a> {
1042        features: &'a [f64],
1043        baseline_estimate: u64,
1044    }
1045
1046    /// JSON response body. Extra fields are ignored.
1047    #[derive(Deserialize)]
1048    struct InferResponse {
1049        estimate: u64,
1050    }
1051
1052    /// HTTP-backed LLM-pluggable corrector.
1053    ///
1054    /// Holds a tiny client config and a base URL. The `ureq` agent is
1055    /// constructed per-call: the per-estimate cost is dominated by LLM
1056    /// inference (hundreds of milliseconds), not agent allocation, and
1057    /// per-call agents keep the struct cheaply `Send + Sync` without
1058    /// interior mutability.
1059    pub struct LlmHttpCorrector {
1060        options: LlmHttpOptions,
1061    }
1062
1063    impl LlmHttpCorrector {
1064        /// Build a corrector from explicit options. The `timeout_ms`
1065        /// value is saturated to [`MAX_TIMEOUT_MS`] so misconfigured
1066        /// callers cannot stall the optimizer for longer than that.
1067        pub fn new(mut options: LlmHttpOptions) -> Self {
1068            if options.timeout_ms > MAX_TIMEOUT_MS {
1069                options.timeout_ms = MAX_TIMEOUT_MS;
1070            }
1071            super::warn_if_remote_plaintext_http(&options.base_url, "llm_http");
1072            Self { options }
1073        }
1074
1075        /// Convenience constructor: default options with the supplied
1076        /// URL. Useful for ad-hoc bench / smoke clients.
1077        pub fn with_url(base_url: impl Into<String>) -> Self {
1078            Self::new(LlmHttpOptions {
1079                base_url: base_url.into(),
1080                ..LlmHttpOptions::default()
1081            })
1082        }
1083
1084        /// Configured options (for diagnostics / logging).
1085        pub fn options(&self) -> &LlmHttpOptions {
1086            &self.options
1087        }
1088
1089        /// Attempt one inference call. Returns `None` on any failure
1090        /// (network, parse, non-2xx). The `correct()` trait method wraps
1091        /// this and applies the LpBound clamp.
1092        fn try_infer(&self, features: &CorrectionFeatures) -> Option<u64> {
1093            let feature_vec = features.to_vec();
1094            let payload = InferRequest {
1095                features: &feature_vec,
1096                baseline_estimate: features.baseline_estimate,
1097            };
1098
1099            let timeout = Duration::from_millis(self.options.timeout_ms);
1100            let agent = ureq::AgentBuilder::new()
1101                .timeout_connect(timeout)
1102                .timeout_read(timeout)
1103                .timeout_write(timeout)
1104                .build();
1105
1106            let response = match agent.post(&self.options.base_url).send_json(&payload) {
1107                Ok(r) => r,
1108                Err(err) => {
1109                    log::debug!(
1110                        "llm_http: request to {} failed: {}",
1111                        self.options.base_url,
1112                        err
1113                    );
1114                    return None;
1115                }
1116            };
1117
1118            match response.into_json::<InferResponse>() {
1119                Ok(body) => Some(body.estimate),
1120                Err(err) => {
1121                    log::debug!(
1122                        "llm_http: response from {} failed to parse: {}",
1123                        self.options.base_url,
1124                        err
1125                    );
1126                    None
1127                }
1128            }
1129        }
1130    }
1131
1132    impl Corrector for LlmHttpCorrector {
1133        fn correct(&self, features: &CorrectionFeatures) -> Result<Option<u64>> {
1134            // Safety contract: every failure returns Ok(None), not Err.
1135            // Mirrors `TabPfnHttpCorrector::correct`. On the optimizer's
1136            // hot path a remote LLM going down (rate limit, network
1137            // partition, mis-config) must never surface as a query
1138            // failure.
1139            let Some(raw) = self.try_infer(features) else {
1140                return Ok(None);
1141            };
1142            Ok(Some(saturating_clamp(raw as f64, self.options.ceiling)))
1143        }
1144
1145        fn name(&self) -> &'static str {
1146            "llm-http"
1147        }
1148    }
1149}
1150
1151#[cfg(test)]
1152mod tests {
1153    use super::*;
1154
1155    #[test]
1156    fn identity_returns_baseline() {
1157        let corrector = IdentityCorrector;
1158        let features = CorrectionFeatures {
1159            baseline_estimate: 1234,
1160            ..Default::default()
1161        };
1162        assert_eq!(corrector.correct(&features).unwrap(), Some(1234));
1163        assert_eq!(corrector.name(), "identity");
1164    }
1165
1166    #[test]
1167    fn tabpfn_stub_always_returns_none() {
1168        let corrector = TabPfnStub;
1169        let features = CorrectionFeatures {
1170            baseline_estimate: 9999,
1171            ..Default::default()
1172        };
1173        assert_eq!(
1174            corrector.correct(&features).unwrap(),
1175            None,
1176            "TabPfnStub must always return Ok(None) — it documents the integration point"
1177        );
1178        assert_eq!(corrector.name(), "tabpfn-stub");
1179
1180        // Also exercise an empty feature vector — the stub should still
1181        // return None without inspecting the input.
1182        let empty = CorrectionFeatures::default();
1183        assert_eq!(corrector.correct(&empty).unwrap(), None);
1184    }
1185
1186    #[test]
1187    fn feature_vec_layout_is_stable() {
1188        let f = CorrectionFeatures {
1189            baseline_estimate: 100,
1190            left_input_rows: Some(10),
1191            right_input_rows: None,
1192            left_distinct: Some(7),
1193            right_distinct: None,
1194            predicate_count: 3,
1195            join_depth: 2,
1196        };
1197        let v = f.to_vec();
1198        assert_eq!(v.len(), CorrectionFeatures::FEATURE_LEN);
1199        assert_eq!(v[0], 100.0);
1200        assert_eq!(v[1], 10.0);
1201        assert_eq!(v[2], 0.0); // None → 0
1202        assert_eq!(v[3], 7.0);
1203        assert_eq!(v[4], 0.0);
1204        assert_eq!(v[5], 3.0);
1205        assert_eq!(v[6], 2.0);
1206    }
1207}
1208
1209#[cfg(all(test, feature = "gbt"))]
1210mod gbt_tests {
1211    use super::gbt::{GbtCorrector, GbtOptions};
1212    use super::{CorrectionFeatures, Corrector};
1213    use crate::feedback::Observation;
1214
1215    /// Build N synthetic observations where `actual = est * 2` for a
1216    /// spread of est values. Plenty of signal for the trees to latch on.
1217    fn synthetic_double(n: u64) -> Vec<Observation> {
1218        (1..=n)
1219            .map(|i| Observation {
1220                template_hash: "syn".into(),
1221                plan_fingerprint: "p".into(),
1222                est_rows: i * 10,
1223                actual_rows: i * 10 * 2,
1224                latency_ms: None,
1225            })
1226            .collect()
1227    }
1228
1229    #[test]
1230    fn predicts_roughly_double_when_training_says_double() {
1231        let obs = synthetic_double(200);
1232        let opts = GbtOptions {
1233            learning_rate: 0.3,
1234            max_depth: 4,
1235            num_trees: 50,
1236            ceiling: u64::MAX,
1237            min_leaf_size: 1,
1238        };
1239        let corrector = GbtCorrector::train(&obs, opts).expect("training");
1240
1241        let features = CorrectionFeatures {
1242            baseline_estimate: 500,
1243            ..Default::default()
1244        };
1245        let corrected = corrector
1246            .correct(&features)
1247            .expect("correct")
1248            .expect("Some");
1249        // True target is 1000. Trees won't be exact; require within 25%.
1250        let ratio = corrected as f64 / 1000.0;
1251        assert!(
1252            (0.75..=1.25).contains(&ratio),
1253            "expected ~1000, got {} (ratio {})",
1254            corrected,
1255            ratio
1256        );
1257        assert_eq!(corrector.name(), "gbt");
1258    }
1259
1260    #[test]
1261    fn ceiling_clamps_when_prediction_exceeds_it() {
1262        let obs = synthetic_double(200);
1263        let opts = GbtOptions {
1264            learning_rate: 0.3,
1265            max_depth: 4,
1266            num_trees: 50,
1267            ceiling: 100, // far below 2 × baseline
1268            min_leaf_size: 1,
1269        };
1270        let corrector = GbtCorrector::train(&obs, opts).expect("training");
1271
1272        let features = CorrectionFeatures {
1273            baseline_estimate: 500,
1274            ..Default::default()
1275        };
1276        let corrected = corrector
1277            .correct(&features)
1278            .expect("correct")
1279            .expect("Some");
1280        assert_eq!(corrected, 100, "ceiling must clamp the corrected estimate");
1281        assert_eq!(corrector.ceiling(), 100);
1282    }
1283
1284    #[test]
1285    fn empty_observations_errors() {
1286        match GbtCorrector::train(&[], GbtOptions::default()) {
1287            Ok(_) => panic!("expected error on empty observations"),
1288            Err(e) => assert!(matches!(e, crate::Error::Feedback(_))),
1289        }
1290    }
1291
1292    #[test]
1293    fn all_zero_observations_errors() {
1294        let obs = vec![
1295            Observation {
1296                template_hash: "z".into(),
1297                plan_fingerprint: "p".into(),
1298                est_rows: 0,
1299                actual_rows: 5,
1300                latency_ms: None,
1301            },
1302            Observation {
1303                template_hash: "z".into(),
1304                plan_fingerprint: "p".into(),
1305                est_rows: 5,
1306                actual_rows: 0,
1307                latency_ms: None,
1308            },
1309        ];
1310        match GbtCorrector::train(&obs, GbtOptions::default()) {
1311            Ok(_) => panic!("expected error when all observations are zero"),
1312            Err(e) => assert!(matches!(e, crate::Error::Feedback(_))),
1313        }
1314    }
1315}
1316
1317#[cfg(all(test, feature = "additive_gbt"))]
1318mod additive_tests {
1319    use super::additive::{AdditiveGbtCorrector, AdditiveGbtOptions};
1320    use super::{CorrectionFeatures, Corrector};
1321    use crate::feedback::Observation;
1322
1323    /// Build N synthetic observations where every actual row count is
1324    /// the same constant `target`. An additive model trained on this
1325    /// should regress toward `target` regardless of the input features.
1326    fn synthetic_constant(n: u64, target: u64) -> Vec<Observation> {
1327        (1..=n)
1328            .map(|i| Observation {
1329                template_hash: "syn-add".into(),
1330                plan_fingerprint: "p".into(),
1331                est_rows: i * 10,
1332                actual_rows: target,
1333                latency_ms: None,
1334            })
1335            .collect()
1336    }
1337
1338    #[test]
1339    fn predicts_near_constant_when_training_is_constant() {
1340        let obs = synthetic_constant(200, 1000);
1341        let opts = AdditiveGbtOptions {
1342            learning_rate: 0.3,
1343            max_depth: 4,
1344            num_trees: 50,
1345            ceiling: u64::MAX,
1346            min_leaf_size: 1,
1347        };
1348        let corrector =
1349            AdditiveGbtCorrector::train(&obs, opts).expect("training additive corrector");
1350
1351        let features = CorrectionFeatures {
1352            baseline_estimate: 500,
1353            ..Default::default()
1354        };
1355        let corrected = corrector
1356            .correct(&features)
1357            .expect("correct")
1358            .expect("Some");
1359        assert!(
1360            (800..=1200).contains(&corrected),
1361            "expected ~1000, got {corrected}"
1362        );
1363        assert_eq!(corrector.name(), "additive_gbt");
1364    }
1365
1366    #[test]
1367    fn ceiling_clamps_when_prediction_exceeds_it() {
1368        let obs = synthetic_constant(200, 1000);
1369        let opts = AdditiveGbtOptions {
1370            learning_rate: 0.3,
1371            max_depth: 4,
1372            num_trees: 50,
1373            ceiling: 100, // far below the trained constant
1374            min_leaf_size: 1,
1375        };
1376        let corrector = AdditiveGbtCorrector::train(&obs, opts).expect("training");
1377
1378        let features = CorrectionFeatures {
1379            baseline_estimate: 500,
1380            ..Default::default()
1381        };
1382        let corrected = corrector
1383            .correct(&features)
1384            .expect("correct")
1385            .expect("Some");
1386        assert_eq!(corrected, 100, "ceiling must clamp the additive correction");
1387        assert_eq!(corrector.ceiling(), 100);
1388    }
1389
1390    #[test]
1391    fn corrects_nonzero_even_when_baseline_estimate_is_zero() {
1392        // This is the q=∞ fix proof. The multiplicative GbtCorrector
1393        // would return 0 here (baseline * exp(predicted) = 0 * _ = 0).
1394        // The additive backend must escape that trap.
1395        let obs = synthetic_constant(200, 1000);
1396        let corrector =
1397            AdditiveGbtCorrector::train(&obs, AdditiveGbtOptions::default()).expect("training");
1398
1399        let features = CorrectionFeatures {
1400            baseline_estimate: 0,
1401            ..Default::default()
1402        };
1403        let corrected = corrector
1404            .correct(&features)
1405            .expect("correct")
1406            .expect("Some");
1407        assert!(
1408            corrected > 0,
1409            "additive corrector must return non-zero even when baseline_estimate = 0; got {corrected}"
1410        );
1411    }
1412
1413    #[test]
1414    fn empty_observations_errors() {
1415        match AdditiveGbtCorrector::train(&[], AdditiveGbtOptions::default()) {
1416            Ok(_) => panic!("expected error on empty observations"),
1417            Err(e) => assert!(matches!(e, crate::Error::Feedback(_))),
1418        }
1419    }
1420}
1421
1422#[cfg(all(test, feature = "tabpfn_http"))]
1423mod tabpfn_http_tests {
1424    use super::tabpfn::{TabPfnHttpCorrector, TabPfnHttpOptions};
1425    use super::{CorrectionFeatures, Corrector};
1426
1427    /// Pointing at port 1 on the loopback interface is the canonical
1428    /// "guaranteed-to-refuse-connection" target on Linux/macOS. The
1429    /// safety contract says: any transport failure must surface as
1430    /// `Ok(None)`, never `Err`, never a panic. We verify that here
1431    /// without standing up a real inference server.
1432    #[test]
1433    fn http_failure_returns_none_not_error() {
1434        let corrector = TabPfnHttpCorrector::new(TabPfnHttpOptions {
1435            base_url: "http://127.0.0.1:1/infer".into(),
1436            timeout_ms: 50,
1437            ceiling: u64::MAX,
1438        });
1439        let features = CorrectionFeatures {
1440            baseline_estimate: 1234,
1441            ..Default::default()
1442        };
1443        let result = corrector.correct(&features);
1444        assert!(
1445            result.is_ok(),
1446            "tabpfn-http transport failure must not propagate as Err; got {result:?}"
1447        );
1448        assert_eq!(
1449            result.unwrap(),
1450            None,
1451            "tabpfn-http transport failure must yield Ok(None) so the engine falls back cleanly"
1452        );
1453        assert_eq!(corrector.name(), "tabpfn-http");
1454    }
1455
1456    #[test]
1457    fn malformed_url_returns_none() {
1458        // Not even a valid URL — `ureq` rejects this at request-build
1459        // time, which our error path must absorb the same as any other
1460        // transport failure.
1461        let corrector = TabPfnHttpCorrector::with_url("not a url at all");
1462        let features = CorrectionFeatures::default();
1463        let result = corrector.correct(&features).expect("never Err");
1464        assert_eq!(result, None);
1465    }
1466
1467    #[test]
1468    fn options_default_is_localhost() {
1469        let opts = TabPfnHttpOptions::default();
1470        assert!(opts.base_url.starts_with("http://"));
1471        assert!(opts.timeout_ms > 0);
1472        assert_eq!(opts.ceiling, u64::MAX);
1473    }
1474}
1475
1476#[cfg(all(test, feature = "llm_http"))]
1477mod llm_http_tests {
1478    use super::llm::{
1479        DEFAULT_TIMEOUT_MS, DEFAULT_URL, LlmHttpCorrector, LlmHttpOptions, MAX_TIMEOUT_MS,
1480    };
1481    use super::{CorrectionFeatures, Corrector};
1482    use std::io::{Read, Write};
1483    use std::net::TcpListener;
1484    use std::sync::atomic::{AtomicUsize, Ordering};
1485    use std::sync::{Arc, Mutex};
1486    use std::thread;
1487    use std::time::Duration;
1488
1489    /// Tiny hand-rolled mock HTTP server, one-shot per accept. We avoid
1490    /// pulling `mockito` (not currently a dep) and keep the test binary
1491    /// lean. The server reads the full request, then writes a fixed
1492    /// response. The handler closure decides what to send so the same
1493    /// scaffolding serves both success and parse-error cases.
1494    fn spawn_mock(
1495        responder: impl Fn(usize) -> Vec<u8> + Send + Sync + 'static,
1496        max_requests: usize,
1497    ) -> (String, Arc<AtomicUsize>) {
1498        let listener = TcpListener::bind("127.0.0.1:0").expect("bind loopback");
1499        let port = listener.local_addr().unwrap().port();
1500        let url = format!("http://127.0.0.1:{port}/infer");
1501        let counter = Arc::new(AtomicUsize::new(0));
1502        let counter_thread = Arc::clone(&counter);
1503        let responder = Arc::new(Mutex::new(responder));
1504        thread::spawn(move || {
1505            listener
1506                .set_nonblocking(false)
1507                .expect("blocking mode for mock");
1508            for stream in listener.incoming().take(max_requests) {
1509                let Ok(mut stream) = stream else { continue };
1510                let _ = stream.set_read_timeout(Some(Duration::from_secs(2)));
1511                let _ = stream.set_write_timeout(Some(Duration::from_secs(2)));
1512                // Drain HTTP request: read headers + body. We pull a
1513                // bounded chunk; the bench client sends tiny payloads
1514                // (sub-200 bytes) so this is sufficient for the tests
1515                // and avoids the parsing complexity of a full HTTP
1516                // server.
1517                let mut buf = [0u8; 4096];
1518                let _ = stream.read(&mut buf);
1519                let idx = counter_thread.fetch_add(1, Ordering::SeqCst);
1520                let body = responder.lock().unwrap()(idx);
1521                let header = format!(
1522                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
1523                    body.len()
1524                );
1525                let _ = stream.write_all(header.as_bytes());
1526                let _ = stream.write_all(&body);
1527                let _ = stream.flush();
1528            }
1529        });
1530        (url, counter)
1531    }
1532
1533    /// Pointing at port 1 on the loopback interface is the canonical
1534    /// "guaranteed-to-refuse-connection" target on Linux/macOS. The
1535    /// safety contract says: any transport failure must surface as
1536    /// `Ok(None)`, never `Err`, never a panic.
1537    #[test]
1538    fn http_failure_returns_none_not_error() {
1539        let corrector = LlmHttpCorrector::new(LlmHttpOptions {
1540            base_url: "http://127.0.0.1:1/infer".into(),
1541            timeout_ms: 50,
1542            ceiling: u64::MAX,
1543        });
1544        let features = CorrectionFeatures {
1545            baseline_estimate: 1234,
1546            ..Default::default()
1547        };
1548        let result = corrector.correct(&features);
1549        assert!(
1550            result.is_ok(),
1551            "llm-http transport failure must not propagate as Err; got {result:?}"
1552        );
1553        assert_eq!(
1554            result.unwrap(),
1555            None,
1556            "llm-http transport failure must yield Ok(None) so the engine falls back cleanly"
1557        );
1558        assert_eq!(corrector.name(), "llm-http");
1559    }
1560
1561    #[test]
1562    fn malformed_url_returns_none() {
1563        let corrector = LlmHttpCorrector::with_url("not a url at all");
1564        let features = CorrectionFeatures::default();
1565        let result = corrector.correct(&features).expect("never Err");
1566        assert_eq!(result, None);
1567    }
1568
1569    #[test]
1570    fn options_default_is_localhost_on_llm_port() {
1571        let opts = LlmHttpOptions::default();
1572        assert_eq!(opts.base_url, DEFAULT_URL);
1573        assert!(opts.base_url.contains(":8766"));
1574        assert_eq!(opts.timeout_ms, DEFAULT_TIMEOUT_MS);
1575        assert_eq!(opts.ceiling, u64::MAX);
1576    }
1577
1578    #[test]
1579    fn timeout_is_saturated_to_max() {
1580        let corrector = LlmHttpCorrector::new(LlmHttpOptions {
1581            base_url: "http://127.0.0.1:1/infer".into(),
1582            timeout_ms: MAX_TIMEOUT_MS * 10,
1583            ceiling: u64::MAX,
1584        });
1585        assert_eq!(corrector.options().timeout_ms, MAX_TIMEOUT_MS);
1586    }
1587
1588    #[test]
1589    fn mock_success_returns_clamped_estimate() {
1590        let (url, counter) = spawn_mock(|_| br#"{"estimate": 4242}"#.to_vec(), 2);
1591        let corrector = LlmHttpCorrector::new(LlmHttpOptions {
1592            base_url: url,
1593            // Generous on purpose: these assert the *response* is handled
1594            // correctly, not that it arrives quickly. A shared CI runner can
1595            // take seconds to spawn the mock thread and complete the loopback
1596            // round trip, and a tight bound here turns that into a flake.
1597            timeout_ms: 15_000,
1598            ceiling: 1_000_000,
1599        });
1600        let features = CorrectionFeatures {
1601            baseline_estimate: 1_000,
1602            ..Default::default()
1603        };
1604        let result = corrector.correct(&features).expect("ok");
1605        assert_eq!(result, Some(4242));
1606        assert!(counter.load(Ordering::SeqCst) >= 1);
1607    }
1608
1609    #[test]
1610    fn mock_clamps_to_ceiling() {
1611        let (url, _counter) = spawn_mock(|_| br#"{"estimate": 99999999}"#.to_vec(), 2);
1612        let corrector = LlmHttpCorrector::new(LlmHttpOptions {
1613            base_url: url,
1614            // Generous on purpose: these assert the *response* is handled
1615            // correctly, not that it arrives quickly. A shared CI runner can
1616            // take seconds to spawn the mock thread and complete the loopback
1617            // round trip, and a tight bound here turns that into a flake.
1618            timeout_ms: 15_000,
1619            ceiling: 500,
1620        });
1621        let result = corrector
1622            .correct(&CorrectionFeatures::default())
1623            .expect("ok");
1624        assert_eq!(result, Some(500));
1625    }
1626
1627    #[test]
1628    fn mock_parse_error_returns_none() {
1629        let (url, _counter) = spawn_mock(|_| b"not json at all".to_vec(), 2);
1630        let corrector = LlmHttpCorrector::with_url(url);
1631        let result = corrector
1632            .correct(&CorrectionFeatures::default())
1633            .expect("ok");
1634        assert_eq!(result, None);
1635    }
1636}