mesh/mesh-rs/crates/mesh-router/src/lib.rs

@@ -14,36 +14,145 @@
 
 use mesh_schema::{BackendId, Fidelity, RouteCandidate, RouteDecision, ServeRequest};
 
+/// How a backend's energy scales with the context it must attend over.
+///
+/// # Why a scalar is not enough
+///
+/// The router previously costed every backend as `uj_per_token × tokens`: one
+/// scalar, linear in token count, blind to how much context precedes each
+/// token. For a backend whose per-token cost genuinely does not depend on the
+/// prefix, that is correct and remains expressible ([`CostProfile::linear`]).
+///
+/// For one that re-reads its context per token it is not, and the error has a
+/// **direction**. A linear model applied to a super-linear backend
+/// *under*-estimates, always. At a gate, under-estimating means admitting work
+/// the budget cannot cover, so the request is accepted and then overruns its
+/// ceiling — which is the one failure a joule ceiling exists to prevent. A
+/// ceiling checked against a model that cannot represent the expensive case is
+/// not a ceiling; it is a formality that happens to pass.
+///
+/// # Why no profile ships as a default
+///
+/// The context term is a property of a *particular backend on particular
+/// hardware*, and must be calibrated by measurement there. Published figures
+/// are not a substitute: the study that motivated this change measured one
+/// model per attention condition, so architecture is confounded with model
+/// family and parameter count, and its slopes cannot be lifted into a constant
+/// without inheriting that confound. So this type carries the *shape* of the
+/// cost and the operator supplies the numbers — the same reason
+/// `BandingScale` ships no canonical scale.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct CostProfile {
+    /// Fixed microjoules to serve the request at all, independent of size.
+    pub base_uj: u64,
+    /// Microjoules per token, independent of context length.
+    pub uj_per_token: u64,
+    /// Additional microjoules per token, per 1024 tokens of preceding context.
+    ///
+    /// Zero reproduces the historical scalar model exactly. A non-zero value
+    /// says this backend gets more expensive per token as the prefix grows,
+    /// which is what makes a long-context request refusable *before* it runs.
+    pub uj_per_token_per_kctx: u64,
+}
+
+/// Tokens of context per unit of [`CostProfile::uj_per_token_per_kctx`].
+const CTX_UNIT: u64 = 1024;
+
+impl CostProfile {
+    /// The historical model: a flat per-token cost with no context term.
+    ///
+    /// Retained so existing callers are unchanged and the scalar case stays
+    /// honestly expressible — some backends really are flat.
+    pub fn linear(uj_per_token: u64) -> Self {
+        CostProfile { base_uj: 0, uj_per_token, uj_per_token_per_kctx: 0 }
+    }
+
+    /// A profile whose per-token cost grows with the context length.
+    pub fn context_scaled(base_uj: u64, uj_per_token: u64, uj_per_token_per_kctx: u64) -> Self {
+        CostProfile { base_uj, uj_per_token, uj_per_token_per_kctx }
+    }
+
+    /// Whether this profile can represent a context-dependent cost at all.
+    ///
+    /// A router gating long-context work against flat profiles is deciding with
+    /// a model that cannot express the case it is meant to catch.
+    pub fn is_context_aware(&self) -> bool {
+        self.uj_per_token_per_kctx > 0
+    }
+
+    /// Estimated microjoules for `tokens` tokens over `context_tokens` of
+    /// preceding context.
+    ///
+    /// Saturating throughout: an estimate that wrapped would come back small,
+    /// and a small estimate is precisely what admits an unaffordable request.
+    pub fn estimate_uj(&self, tokens: u64, context_tokens: u64) -> u64 {
+        let flat = self.uj_per_token.saturating_mul(tokens);
+        let ctx_units = context_tokens / CTX_UNIT;
+        let ctx = self
+            .uj_per_token_per_kctx
+            .saturating_mul(tokens)
+            .saturating_mul(ctx_units);
+        self.base_uj.saturating_add(flat).saturating_add(ctx)
+    }
+
+    /// The profile as it is sealed into a receipt, so a route decision can be
+    /// re-derived rather than trusted.
+    pub fn wire(&self) -> String {
+        format!(
+            "base={};per_token={};per_token_per_kctx={}",
+            self.base_uj, self.uj_per_token, self.uj_per_token_per_kctx
+        )
+    }
+}
+
 /// A backend the router can route to: its identity, the recipe it serves, the
-/// fidelity it delivers, and its modeled per-token energy cost.
+/// fidelity it delivers, and its modeled energy cost.
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct BackendSpec {
     pub id: BackendId,
     pub recipe: String,
     pub fidelity: Fidelity,
-    pub uj_per_token: u64,
+    /// How this backend's cost scales. See [`CostProfile`].
+    pub cost: CostProfile,
 }
 
 impl BackendSpec {
+    /// A backend with a flat per-token cost — the historical constructor,
+    /// unchanged in behaviour.
     pub fn new(
         id: impl Into<String>,
         recipe: impl Into<String>,
         fidelity: Fidelity,
         uj_per_token: u64,
     ) -> Self {
-        Self {
+        Self::with_cost(id, recipe, fidelity, CostProfile::linear(uj_per_token))
-            id: BackendId::new(id),
-            recipe: recipe.into(),
-            fidelity,
-            uj_per_token,
-        }
     }
 
-    /// Estimated total microjoules to serve `req` on this backend: per-token
+    /// A backend with an explicit cost profile.
-    /// cost over input + output tokens.
+    pub fn with_cost(
+        id: impl Into<String>,
+        recipe: impl Into<String>,
+        fidelity: Fidelity,
+        cost: CostProfile,
+    ) -> Self {
+        Self { id: BackendId::new(id), recipe: recipe.into(), fidelity, cost }
+    }
+
+    /// The flat per-token term, for callers that only ever set one.
+    pub fn uj_per_token(&self) -> u64 {
+        self.cost.uj_per_token
+    }
+
+    /// Estimated total microjoules to serve `req` on this backend.
+    ///
+    /// The context term is charged against `prompt_tokens`: the prefix every
+    /// generated token attends over. Output tokens extend that prefix as they
+    /// are produced, so this remains a lower bound on a growing-context
+    /// backend — deliberately, since the alternative is to model the growth and
+    /// then defend a second uncalibrated constant.
     fn estimate(&self, req: &ServeRequest) -> u64 {
         let tokens = req.prompt_tokens as u64 + req.max_tokens as u64;
-        self.uj_per_token.saturating_mul(tokens)
+        self.cost.estimate_uj(tokens, req.prompt_tokens as u64)
     }
 }
 
@@ -168,6 +277,116 @@ mod tests {
         r
     }
 
+    fn req_ctx(floor: Fidelity, ceiling: u64, prompt_tokens: u32, max_tokens: u32) -> ServeRequest {
+        ServeRequest {
+            prompt_hash: ContentId::of_bytes(b"q"),
+            prompt_tokens,
+            max_tokens,
+            fidelity_floor: floor,
+            joule_ceiling_uj: ceiling,
+        }
+    }
+
+    /// The historical scalar model must be reproduced exactly, so existing
+    /// backends and their receipts do not shift under this change.
+    #[test]
+    fn a_linear_profile_is_the_old_model_exactly() {
+        let p = CostProfile::linear(7);
+        assert_eq!(p.estimate_uj(100, 0), 700);
+        // Context is irrelevant to a flat backend, however long it gets.
+        assert_eq!(p.estimate_uj(100, 1_000_000), 700);
+        assert!(!p.is_context_aware());
+
+        // And via the untouched constructor.
+        let b = BackendSpec::new("b", "r", Fidelity::F1, 7);
+        assert_eq!(b.estimate(&req_ctx(Fidelity::F1, u64::MAX, 10, 90)), 700);
+        assert_eq!(b.uj_per_token(), 7);
+    }
+
+    /// The defect this change exists for: a flat model applied to a
+    /// context-scaling backend under-estimates, and the error grows with the
+    /// context. Under-estimating at a gate means *admitting* work the ceiling
+    /// cannot cover.
+    #[test]
+    fn a_flat_model_under_estimates_a_context_scaling_backend() {
+        let flat = CostProfile::linear(10);
+        let scaled = CostProfile::context_scaled(0, 10, 2);
+
+        // Short prompt: the two models agree, which is why the defect hides.
+        assert_eq!(flat.estimate_uj(1_000, 512), scaled.estimate_uj(1_000, 512));
+
+        // Long prompt: they diverge, and always in the same direction.
+        let ctx = 16_384;
+        let f = flat.estimate_uj(1_000, ctx);
+        let s = scaled.estimate_uj(1_000, ctx);
+        assert!(s > f, "the context term must cost more, not less");
+        // 16_384 / 1024 = 16 context units -> 10*1000 + 2*1000*16.
+        assert_eq!(s, 10_000 + 32_000);
+        assert_eq!(f, 10_000);
+    }
+
+    /// The capability the docs imply and the scalar model could not deliver:
+    /// refuse a long-context route *before* it runs.
+    #[test]
+    fn a_long_context_route_is_refused_before_it_burns_the_budget() {
+        let ceiling = 20_000u64;
+        // 100 output tokens over a 1_500-token prompt.
+        let r = req_ctx(Fidelity::F1, ceiling, 1_500, 100);
+
+        // Described flat, the same backend is costed at 10 * 1_600 = 16_000 and
+        // is admitted.
+        let mut blind = Router::new();
+        blind.register(BackendSpec::new("b", "r", Fidelity::F1, 10));
+        let blind_est = blind.spec(&BackendId::new("b")).unwrap().estimate(&r);
+        assert_eq!(blind_est, 16_000);
+        assert!(blind_est <= ceiling, "the flat model admits it");
+        assert!(blind.route(&r).chosen.is_some(), "and the router therefore routes it");
+
+        // Described with a calibrated context term, the same request costs more
+        // than the ceiling and is refused before anything runs.
+        let mut aware = Router::new();
+        aware.register(BackendSpec::with_cost(
+            "b",
+            "r",
+            Fidelity::F1,
+            CostProfile::context_scaled(0, 10, 200),
+        ));
+        let aware_est = aware.spec(&BackendId::new("b")).unwrap().estimate(&r);
+        assert!(aware_est > ceiling, "the calibrated model refuses it ({aware_est})");
+        assert!(
+            aware.route(&r).chosen.is_none(),
+            "the router must not route what it cannot afford"
+        );
+    }
+
+    /// A profile that cannot express context is detectable, so a deployment
+    /// gating long-context work on flat profiles can notice it is deciding with
+    /// a model that cannot represent the case.
+    #[test]
+    fn context_awareness_is_inspectable() {
+        assert!(!CostProfile::linear(10).is_context_aware());
+        assert!(CostProfile::context_scaled(0, 10, 1).is_context_aware());
+        assert!(!CostProfile::context_scaled(500, 10, 0).is_context_aware());
+    }
+
+    /// An estimate must never wrap: a wrapped estimate comes back small, and a
+    /// small estimate is exactly what admits an unaffordable request.
+    #[test]
+    fn an_overflowing_estimate_saturates_rather_than_wrapping() {
+        let p = CostProfile::context_scaled(u64::MAX, u64::MAX, u64::MAX);
+        assert_eq!(p.estimate_uj(u64::MAX, u64::MAX), u64::MAX);
+        // And a saturated estimate exceeds any real ceiling, so it refuses.
+        assert!(p.estimate_uj(1_000_000, 1_000_000) > 1_000_000_000);
+    }
+
+    /// The profile is sealable, so a route decision can be re-derived from the
+    /// receipt rather than trusted.
+    #[test]
+    fn the_profile_is_carried_on_the_wire() {
+        let w = CostProfile::context_scaled(5, 10, 2).wire();
+        assert_eq!(w, "base=5;per_token=10;per_token_per_kctx=2");
+    }
+
     fn req(floor: Fidelity, ceiling: u64, max_tokens: u32) -> ServeRequest {
         ServeRequest {
             prompt_hash: ContentId::of_bytes(b"q"),