Skip to main content

capsule_runtime/
model_limits.rs

1//! Each driver choice's per-turn output cap, context window, input ceiling and input budget,
2//! resolved once at launch.
3//!
4//! A `transport: http` capsule names a model and the driver that serves it. The driver's bundled
5//! `murmur.yaml` declares, in its `models:` table, each model's `max_output_tokens`,
6//! `context_window` and optionally `max_input_tokens` (read by
7//! [`murmur_artifact::parse_model_limits`]). For every driver choice, the primary and each
8//! `inference.alternates` entry:
9//!
10//! | Limit | `"manifest"` | else `"driver"` | else |
11//! |---|---|---|---|
12//! | output cap | `inference.max_tokens` | the table's `max_output_tokens` | refused, `E-RUN-052` |
13//! | context window | `context.max_tokens` | the table's `context_window` | none, `"none"` |
14//! | input ceiling | — | the table's `max_input_tokens` | none, `"none"` |
15//!
16//! The manifest's two keys are session-wide: each applies to every choice as written, never
17//! clamped to a declared ceiling. A model id matches a table entry by exact string equality. A
18//! choice whose model the table does not list draws `W-RUN-009` unless the manifest sets both
19//! keys.
20//!
21//! Every inference call has one input budget, `min(input ceiling, context window − output cap)`
22//! with the cap the call sends ([`ResolvedModelLimits::input_budget_for`]). Compaction's
23//! threshold and the seed budget are fractions of it. With no window there is no budget: for
24//! that choice compaction does not run and no seed budget applies. A choice whose budget is `0`
25//! is refused, `E-RUN-055`.
26//!
27//! `inference.compaction.max_tokens` is checked here too: above the primary driver's ceiling for
28//! the compaction model it is refused, `E-RUN-054`; at or above the compaction model's window it
29//! leaves compaction calls no input budget, `E-RUN-055`.
30//!
31//! Staging and `mur doctor` both resolve through [`resolve_model_limits`], so the two never
32//! disagree on a number or on the warning's text. `transport: process` resolves nothing: the
33//! harness owns its model and its output.
34
35use std::collections::BTreeMap;
36
37use murmur_artifact::{
38    ContextConfig, InferenceConfig, ModelLimits, PRIMARY_DRIVER_CHOICE, W_RUN_009,
39};
40use serde::Serialize;
41
42use crate::inference_import::HookInferencePurpose;
43
44/// Where a limit in force came from, as `session_start` names it.
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
46#[serde(rename_all = "lowercase")]
47pub enum LimitSource {
48    /// The manifest's `inference.max_tokens` or `context.max_tokens`.
49    Manifest,
50    /// The driver's `models:` table entry for the choice's model.
51    Driver,
52    /// Neither: only ever a context window, which the choice then does not have.
53    None,
54}
55
56impl LimitSource {
57    /// The value as written in `trace.jsonl` and printed by `mur doctor`.
58    #[must_use]
59    pub fn wire_name(self) -> &'static str {
60        match self {
61            Self::Manifest => "manifest",
62            Self::Driver => "driver",
63            Self::None => "none",
64        }
65    }
66}
67
68impl std::fmt::Display for LimitSource {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        f.write_str(self.wire_name())
71    }
72}
73
74/// Which term of `min(input ceiling, context window − output cap)` an input budget is.
75#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
76#[serde(rename_all = "snake_case")]
77pub enum BudgetBasis {
78    /// The driver's `max_input_tokens`, at most the window less the cap (a tie goes here).
79    MaxInputTokens,
80    /// The context window less the output cap, below any ceiling.
81    ContextWindowLessOutputCap,
82    /// The context window itself, where no output cap is known (`transport: process`).
83    ContextWindow,
84    /// No window, so no budget.
85    None,
86}
87
88impl BudgetBasis {
89    /// The value as written in `trace.jsonl`, `session-status` and printed by `mur doctor`.
90    #[must_use]
91    pub fn wire_name(self) -> &'static str {
92        match self {
93            Self::MaxInputTokens => "max_input_tokens",
94            Self::ContextWindowLessOutputCap => "context_window_less_output_cap",
95            Self::ContextWindow => "context_window",
96            Self::None => "none",
97        }
98    }
99}
100
101/// The most input tokens one inference call may carry, and which limit set it.
102#[derive(Debug, Clone, Copy, PartialEq, Eq)]
103pub struct InputBudget {
104    pub tokens: u32,
105    pub basis: BudgetBasis,
106}
107
108impl InputBudget {
109    /// `min(ceiling, window − cap)`, saturating at `0`; the ceiling wins a tie.
110    #[must_use]
111    pub fn of(window: u32, cap: u32, ceiling: Option<u32>) -> Self {
112        let rest = window.saturating_sub(cap);
113        match ceiling {
114            Some(ceiling) if ceiling <= rest => Self {
115                tokens: ceiling,
116                basis: BudgetBasis::MaxInputTokens,
117            },
118            _ => Self {
119                tokens: rest,
120                basis: BudgetBasis::ContextWindowLessOutputCap,
121            },
122        }
123    }
124}
125
126/// `n`, or `none` for an absent limit, as the refusals, `mur doctor` and `MURMUR.md` print it.
127#[must_use]
128pub fn limit_or_none(limit: Option<u32>) -> String {
129    limit.map_or_else(|| "none".to_string(), |limit| limit.to_string())
130}
131
132/// One driver choice's output cap, context window, input ceiling and input budget in force, and
133/// where each came from. The `model_limits` object of `session_start` and of each
134/// `inference_choices` element.
135#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
136pub struct ResolvedModelLimits {
137    /// The `max_tokens` every agent-loop call under this choice sends, unless a controller has
138    /// changed `inference.max_tokens`.
139    pub max_output_tokens: u32,
140    /// [`LimitSource::Manifest`] or [`LimitSource::Driver`].
141    pub max_output_tokens_source: LimitSource,
142    /// The window the input budget is cut from; `None` when neither the manifest nor the driver
143    /// gives one.
144    pub context_window: Option<u32>,
145    pub context_window_source: LimitSource,
146    /// The most input tokens the driver says the provider accepts for the model.
147    pub max_input_tokens: Option<u32>,
148    /// [`LimitSource::Driver`] when the table entry declares one, else [`LimitSource::None`].
149    pub max_input_tokens_source: LimitSource,
150    /// [`Self::input_budget_for`] `max_output_tokens`: what compaction and the seed budget
151    /// measure against until a controller changes the cap. `None` with no window.
152    pub input_budget: Option<u32>,
153    pub input_budget_basis: BudgetBasis,
154}
155
156impl ResolvedModelLimits {
157    /// The input budget of a call under this choice that sends `cap` as its output cap; `None`
158    /// with no window.
159    #[must_use]
160    pub fn input_budget_for(&self, cap: u32) -> Option<InputBudget> {
161        self.context_window
162            .map(|window| InputBudget::of(window, cap, self.max_input_tokens))
163    }
164}
165
166/// One driver choice's resolution.
167#[derive(Debug, Clone, PartialEq, Eq)]
168pub struct ChoiceModelLimits {
169    /// [`PRIMARY_DRIVER_CHOICE`] or the alternate's `name`.
170    pub choice: String,
171    pub model: String,
172    /// The driver artifact as `name@version`.
173    pub driver: String,
174    pub limits: ResolvedModelLimits,
175    /// `W-RUN-009`'s text, code and link not included, when the driver's table does not list the
176    /// model and the manifest does not set both keys.
177    pub warning: Option<String>,
178    /// The table entry's `vision`; `None` when the entry does not say or the table does not list
179    /// the model.
180    pub vision: Option<bool>,
181}
182
183/// A driver choice's driver, as resolution reads it.
184#[derive(Debug, Clone, Copy)]
185pub struct DriverManifest<'a> {
186    /// The driver artifact as `name@version`, for the warning and the refusal.
187    pub label: &'a str,
188    /// The driver's bundled `murmur.yaml`.
189    pub manifest_yaml: &'a str,
190}
191
192/// Why a launch is refused: `E-RUN-052`, `E-RUN-054` for
193/// [`ModelLimitsError::CompactionCapAboveCeiling`], or `E-RUN-055` for
194/// [`ModelLimitsError::NoInputBudget`] and [`ModelLimitsError::CompactionNoInputBudget`].
195#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
196pub enum ModelLimitsError {
197    /// A choice has no output cap: its driver's table does not list its model and the manifest
198    /// sets no `inference.max_tokens`.
199    #[error(
200        "driver choice '{choice}' has no output cap: driver {driver} declares no limits for model \
201         '{model}' and the manifest sets no inference.max_tokens"
202    )]
203    NoOutputCap {
204        choice: String,
205        model: String,
206        driver: String,
207    },
208    /// A choice's driver ships a `models:` table that does not parse.
209    #[error("driver {driver} ships a malformed models: table: {reason}")]
210    MalformedTable { driver: String, reason: String },
211    /// `inference.compaction.max_tokens` exceeds the primary driver's `max_output_tokens` for the
212    /// compaction model.
213    #[error(
214        "inference.compaction.max_tokens ({max_tokens}) is above the output ceiling of \
215         {ceiling} that driver {driver} declares for compaction model '{model}'"
216    )]
217    CompactionCapAboveCeiling {
218        max_tokens: u32,
219        ceiling: u32,
220        model: String,
221        driver: String,
222    },
223    /// A choice's output cap is at or above its context window, so its calls have no input
224    /// budget.
225    #[error(
226        "driver choice '{choice}' has no input budget: its context window of {window} \
227         ({window_source}) less its output cap of {cap} ({cap_source}) leaves no input tokens; \
228         input ceiling {ceiling_text} (model '{model}', driver {driver})",
229        ceiling_text = limit_or_none(*.ceiling)
230    )]
231    NoInputBudget {
232        choice: String,
233        model: String,
234        driver: String,
235        window: u32,
236        window_source: LimitSource,
237        cap: u32,
238        cap_source: LimitSource,
239        ceiling: Option<u32>,
240    },
241    /// `inference.compaction.max_tokens` is at or above the compaction model's context window,
242    /// so compaction calls have no input budget.
243    #[error(
244        "compaction calls have no input budget: the context window of {window} \
245         ({window_source}) for compaction model '{model}' less inference.compaction.max_tokens \
246         of {cap} leaves no input tokens; input ceiling {ceiling_text} (driver {driver})",
247        ceiling_text = limit_or_none(*.ceiling)
248    )]
249    CompactionNoInputBudget {
250        model: String,
251        driver: String,
252        window: u32,
253        window_source: LimitSource,
254        cap: u32,
255        ceiling: Option<u32>,
256    },
257}
258
259/// Resolves the output cap, context window, input ceiling and input budget of every driver
260/// choice `inference` declares, the primary first and the alternates in manifest order. Empty
261/// under `transport: process`.
262///
263/// `driver_manifest` takes a choice's driver as `inference` names it (its tool name) to that
264/// driver's label and bundled manifest. Each driver's table is parsed once.
265pub fn resolve_model_limits<'a>(
266    inference: &InferenceConfig,
267    context: Option<&ContextConfig>,
268    mut driver_manifest: impl FnMut(&str) -> DriverManifest<'a>,
269) -> Result<Vec<ChoiceModelLimits>, ModelLimitsError> {
270    if inference.transport != "http" {
271        return Ok(Vec::new());
272    }
273    let primary_driver = inference
274        .driver
275        .as_ref()
276        .map(|driver| driver.artifact.as_str())
277        .unwrap_or_default();
278    let choices = std::iter::once((PRIMARY_DRIVER_CHOICE, primary_driver, &inference.model)).chain(
279        inference.alternates.iter().map(|alternate| {
280            (
281                alternate.name.as_str(),
282                alternate.driver.as_str(),
283                &alternate.model,
284            )
285        }),
286    );
287    let cap_override = inference.max_tokens;
288    let window_override = context.and_then(|context| context.max_tokens);
289    let mut tables: BTreeMap<&str, BTreeMap<String, ModelLimits>> = BTreeMap::new();
290    let mut resolved = Vec::new();
291    for (choice, driver, model) in choices {
292        let manifest = driver_manifest(driver);
293        let table =
294            match tables.get(driver) {
295                Some(table) => table,
296                None => {
297                    let table = murmur_artifact::parse_model_limits(manifest.manifest_yaml)
298                        .map_err(|err| ModelLimitsError::MalformedTable {
299                            driver: manifest.label.to_string(),
300                            reason: err.to_string(),
301                        })?;
302                    tables.entry(driver).or_insert(table)
303                }
304            };
305        let declared = table.get(model.as_str());
306        let (max_output_tokens, max_output_tokens_source) = match (cap_override, declared) {
307            (Some(cap), _) => (cap, LimitSource::Manifest),
308            (None, Some(declared)) => (declared.max_output_tokens, LimitSource::Driver),
309            (None, None) => {
310                return Err(ModelLimitsError::NoOutputCap {
311                    choice: choice.to_string(),
312                    model: model.clone(),
313                    driver: manifest.label.to_string(),
314                })
315            }
316        };
317        let (context_window, context_window_source) = match (window_override, declared) {
318            (Some(window), _) => (Some(window), LimitSource::Manifest),
319            (None, Some(declared)) => (Some(declared.context_window), LimitSource::Driver),
320            (None, None) => (None, LimitSource::None),
321        };
322        // Undeclared with both keys set needs nothing from the table, so there is nothing to say.
323        let warning = (declared.is_none() && context_window.is_none()).then(|| {
324            format!(
325                "driver {} declares no limits for model '{model}' (driver choice '{choice}'), so \
326                 the session uses inference.max_tokens ({max_output_tokens}) as its output cap \
327                 and knows no context window: compaction does not run and no seed budget \
328                 applies; set context.max_tokens to give it one",
329                manifest.label
330            )
331        });
332        let max_input_tokens = declared.and_then(|declared| declared.max_input_tokens);
333        let mut limits = ResolvedModelLimits {
334            max_output_tokens,
335            max_output_tokens_source,
336            context_window,
337            context_window_source,
338            max_input_tokens,
339            max_input_tokens_source: if max_input_tokens.is_some() {
340                LimitSource::Driver
341            } else {
342                LimitSource::None
343            },
344            input_budget: None,
345            input_budget_basis: BudgetBasis::None,
346        };
347        if let Some(budget) = limits.input_budget_for(max_output_tokens) {
348            if budget.tokens == 0 {
349                return Err(ModelLimitsError::NoInputBudget {
350                    choice: choice.to_string(),
351                    model: model.clone(),
352                    driver: manifest.label.to_string(),
353                    window: context_window.unwrap_or_default(),
354                    window_source: context_window_source,
355                    cap: max_output_tokens,
356                    cap_source: max_output_tokens_source,
357                    ceiling: max_input_tokens,
358                });
359            }
360            limits.input_budget = Some(budget.tokens);
361            limits.input_budget_basis = budget.basis;
362        }
363        resolved.push(ChoiceModelLimits {
364            choice: choice.to_string(),
365            model: model.clone(),
366            driver: manifest.label.to_string(),
367            limits,
368            warning,
369            vision: declared.and_then(|declared| declared.vision),
370        });
371    }
372    check_compaction_cap(inference, primary_driver, &tables, &resolved)?;
373    check_compaction_budget(
374        inference,
375        window_override,
376        primary_driver,
377        &tables,
378        &resolved,
379    )?;
380    Ok(resolved)
381}
382
383/// Refuses an `inference.compaction.max_tokens` above the primary driver's ceiling for the
384/// compaction model: the primary driver serves every compaction call. A model its table does not
385/// list has no ceiling to check against.
386fn check_compaction_cap(
387    inference: &InferenceConfig,
388    primary_driver: &str,
389    tables: &BTreeMap<&str, BTreeMap<String, ModelLimits>>,
390    resolved: &[ChoiceModelLimits],
391) -> Result<(), ModelLimitsError> {
392    let Some(compaction) = inference.compaction.as_ref() else {
393        return Ok(());
394    };
395    let Some(max_tokens) = compaction.max_tokens else {
396        return Ok(());
397    };
398    let model = compaction.model.as_deref().unwrap_or(&inference.model);
399    let Some(declared) = tables
400        .get(primary_driver)
401        .and_then(|table| table.get(model))
402    else {
403        return Ok(());
404    };
405    if max_tokens <= declared.max_output_tokens {
406        return Ok(());
407    }
408    Err(ModelLimitsError::CompactionCapAboveCeiling {
409        max_tokens,
410        ceiling: declared.max_output_tokens,
411        model: model.to_string(),
412        driver: resolved
413            .first()
414            .map(|primary| primary.driver.clone())
415            .unwrap_or_default(),
416    })
417}
418
419/// Refuses an `inference.compaction.max_tokens` at or above the compaction model's window:
420/// `context.max_tokens`, else the primary driver's table entry for the model. A model with
421/// neither has no window to check against.
422fn check_compaction_budget(
423    inference: &InferenceConfig,
424    window_override: Option<u32>,
425    primary_driver: &str,
426    tables: &BTreeMap<&str, BTreeMap<String, ModelLimits>>,
427    resolved: &[ChoiceModelLimits],
428) -> Result<(), ModelLimitsError> {
429    let Some(cap) = inference
430        .compaction
431        .as_ref()
432        .and_then(|compaction| compaction.max_tokens)
433    else {
434        return Ok(());
435    };
436    let model = inference
437        .compaction
438        .as_ref()
439        .and_then(|compaction| compaction.model.as_deref())
440        .unwrap_or(&inference.model);
441    let declared = tables
442        .get(primary_driver)
443        .and_then(|table| table.get(model));
444    let (window, window_source) = match (window_override, declared) {
445        (Some(window), _) => (window, LimitSource::Manifest),
446        (None, Some(declared)) => (declared.context_window, LimitSource::Driver),
447        (None, None) => return Ok(()),
448    };
449    let ceiling = declared.and_then(|declared| declared.max_input_tokens);
450    if InputBudget::of(window, cap, ceiling).tokens > 0 {
451        return Ok(());
452    }
453    Err(ModelLimitsError::CompactionNoInputBudget {
454        model: model.to_string(),
455        driver: resolved
456            .first()
457            .map(|primary| primary.driver.clone())
458            .unwrap_or_default(),
459        window,
460        window_source,
461        cap,
462        ceiling,
463    })
464}
465
466/// `W-RUN-009` as staging and `mur doctor` print it, `message` being
467/// [`ChoiceModelLimits::warning`].
468#[must_use]
469pub fn model_limits_warning_line(message: &str) -> String {
470    format!(
471        "warning[{W_RUN_009}]: {message} ({})",
472        murmur_artifact::runtime_warning_link(W_RUN_009)
473    )
474}
475
476/// The output cap a hook's `run-inference` sends: `inference.compaction.max_tokens` while the hook
477/// serves `on-compaction` and the key is set, else the primary driver's table entry for the model
478/// it calls, else the primary choice's resolved cap. `inference.max_tokens` is an agent-turn cap,
479/// so it reaches a hook completion only through that fallback: a small one never truncates a
480/// compaction summary for a model the table knows. A compaction call is capped by its own key
481/// instead, which the operator sets for that call alone.
482///
483/// The call's input budget is cut the same way as an agent-loop call's
484/// ([`InputBudget::of`]), from `context.max_tokens`, else the table entry's window.
485#[derive(Debug, Clone, PartialEq, Eq, Default)]
486pub(crate) struct HookOutputCaps {
487    /// The primary driver's `models:` table.
488    table: BTreeMap<String, ModelLimits>,
489    /// The primary choice's resolved output cap.
490    fallback: u32,
491    /// `inference.compaction.max_tokens`.
492    compaction: Option<u32>,
493    /// `context.max_tokens`.
494    window: Option<u32>,
495}
496
497impl HookOutputCaps {
498    pub(crate) fn new(
499        table: BTreeMap<String, ModelLimits>,
500        fallback: u32,
501        compaction: Option<u32>,
502        window: Option<u32>,
503    ) -> Self {
504        Self {
505            table,
506            fallback,
507            compaction,
508            window,
509        }
510    }
511
512    /// The input budget of a completion on `model` that sends `cap` ([`Self::for_call`]'s
513    /// result): `None` with no window, `0` when the cap fills a window the launch did not check.
514    pub(crate) fn input_budget(&self, model: &str, cap: u32) -> Option<u32> {
515        let declared = self.table.get(model);
516        let window = self
517            .window
518            .or_else(|| declared.map(|limits| limits.context_window))?;
519        Some(
520            InputBudget::of(
521                window,
522                cap,
523                declared.and_then(|limits| limits.max_input_tokens),
524            )
525            .tokens,
526        )
527    }
528
529    /// The cap a completion on `model` made for `purpose` sends and admits spend against.
530    pub(crate) fn for_call(&self, model: &str, purpose: HookInferencePurpose) -> u32 {
531        match (purpose, self.compaction) {
532            (HookInferencePurpose::Compaction, Some(cap)) => cap,
533            _ => self
534                .table
535                .get(model)
536                .map_or(self.fallback, |limits| limits.max_output_tokens),
537        }
538    }
539}
540
541#[cfg(test)]
542mod tests {
543    use super::*;
544    use murmur_artifact::RuntimeManifest;
545
546    const DRIVER: &str = "name: drv\nversion: 0.1.0\nruntime: driver\nmodels:\n  m:\n    \
547        max_output_tokens: 65536\n    context_window: 400000\n  alt-m:\n    \
548        max_output_tokens: 1000\n    context_window: 2000\n";
549    const BARE_DRIVER: &str = "name: bare\nversion: 0.2.0\nruntime: driver\n";
550
551    /// A capsule on `local/drv` with `model`, the `inference` lines `extra`, and `context`, and
552    /// `local/bare` declared beside it for an alternate to name.
553    fn manifest(model: &str, extra: &str, context: &str) -> RuntimeManifest {
554        let entry = |name: &str, version: &str| {
555            format!(
556                "  - name: {name}\n    version: {version}\n    runtime: driver\n    gateway:\n      \
557                 endpoint: http://127.0.0.1:1\n      keyless: true\n"
558            )
559        };
560        let bare = if extra.contains("local/bare") {
561            entry("local/bare", "0.2.0")
562        } else {
563            String::new()
564        };
565        RuntimeManifest::from_yaml_str(&format!(
566            "name: cap\nversion: 0.0.1\nartifacts:\n{}{bare}inference:\n  transport: http\n  \
567             model: {model}\n  driver:\n    artifact: local/drv\n{extra}{context}",
568            entry("local/drv", "0.1.0")
569        ))
570        .expect("manifest parses")
571    }
572
573    fn resolve(manifest: &RuntimeManifest) -> Result<Vec<ChoiceModelLimits>, ModelLimitsError> {
574        resolve_model_limits(
575            manifest.inference.as_ref().unwrap(),
576            manifest.context.as_ref(),
577            |driver| match driver {
578                "local/bare" => DriverManifest {
579                    label: "local/bare@0.2.0",
580                    manifest_yaml: BARE_DRIVER,
581                },
582                _ => DriverManifest {
583                    label: "local/drv@0.1.0",
584                    manifest_yaml: DRIVER,
585                },
586            },
587        )
588    }
589
590    #[test]
591    fn a_listed_model_takes_both_limits_from_the_driver() {
592        let resolved = resolve(&manifest("m", "", "")).unwrap();
593        assert_eq!(resolved.len(), 1);
594        assert_eq!(resolved[0].choice, PRIMARY_DRIVER_CHOICE);
595        assert_eq!(
596            resolved[0].limits,
597            ResolvedModelLimits {
598                max_output_tokens: 65536,
599                max_output_tokens_source: LimitSource::Driver,
600                context_window: Some(400_000),
601                context_window_source: LimitSource::Driver,
602                max_input_tokens: None,
603                max_input_tokens_source: LimitSource::None,
604                input_budget: Some(334_464),
605                input_budget_basis: BudgetBasis::ContextWindowLessOutputCap,
606            }
607        );
608        assert_eq!(resolved[0].warning, None);
609    }
610
611    #[test]
612    fn each_manifest_key_overrides_its_own_limit_as_written() {
613        let both = resolve(&manifest(
614            "m",
615            "  max_tokens: 4000\n",
616            "context:\n  max_tokens: 5000\n",
617        ))
618        .unwrap();
619        assert_eq!(
620            both[0].limits,
621            ResolvedModelLimits {
622                max_output_tokens: 4000,
623                max_output_tokens_source: LimitSource::Manifest,
624                context_window: Some(5000),
625                context_window_source: LimitSource::Manifest,
626                max_input_tokens: None,
627                max_input_tokens_source: LimitSource::None,
628                input_budget: Some(1000),
629                input_budget_basis: BudgetBasis::ContextWindowLessOutputCap,
630            }
631        );
632        let cap = resolve(&manifest("m", "  max_tokens: 4096\n", "")).unwrap();
633        assert_eq!(
634            cap[0].limits.max_output_tokens_source,
635            LimitSource::Manifest
636        );
637        assert_eq!(cap[0].limits.context_window_source, LimitSource::Driver);
638    }
639
640    #[test]
641    fn an_unlisted_model_with_a_cap_warns_and_has_no_window() {
642        let resolved = resolve(&manifest("M", "  max_tokens: 4096\n", "")).unwrap();
643        assert_eq!(resolved[0].limits.max_output_tokens, 4096);
644        assert_eq!(resolved[0].limits.context_window, None);
645        assert_eq!(resolved[0].limits.context_window_source, LimitSource::None);
646        let warning = resolved[0].warning.as_deref().expect("W-RUN-009");
647        assert!(
648            warning.starts_with(
649                "driver local/drv@0.1.0 declares no limits for model 'M' (driver choice \
650                 'primary'), so the session uses inference.max_tokens (4096) as its output cap"
651            ),
652            "{warning}"
653        );
654        assert!(warning.contains("set context.max_tokens"), "{warning}");
655        let both = resolve(&manifest(
656            "M",
657            "  max_tokens: 4096\n",
658            "context:\n  max_tokens: 9000\n",
659        ))
660        .unwrap();
661        assert_eq!(
662            both[0].warning, None,
663            "both keys set need nothing from the table"
664        );
665    }
666
667    #[test]
668    fn an_unlisted_model_without_a_cap_is_refused() {
669        assert_eq!(
670            resolve(&manifest("m-2026", "", "context:\n  max_tokens: 9000\n")),
671            Err(ModelLimitsError::NoOutputCap {
672                choice: "primary".to_string(),
673                model: "m-2026".to_string(),
674                driver: "local/drv@0.1.0".to_string(),
675            })
676        );
677    }
678
679    #[test]
680    fn every_alternate_resolves_against_its_own_driver() {
681        let alternates = "  alternates:\n    - name: cheap\n      driver:\n        artifact: \
682                          local/drv\n      model: alt-m\n    - name: other\n      driver:\n        \
683                          artifact: local/bare\n      model: x\n";
684        let control = "control:\n  agent_settings: [inference.driver]\n";
685        let refused = resolve(&manifest("m", alternates, control));
686        assert_eq!(
687            refused,
688            Err(ModelLimitsError::NoOutputCap {
689                choice: "other".to_string(),
690                model: "x".to_string(),
691                driver: "local/bare@0.2.0".to_string(),
692            })
693        );
694        let with_cap = format!("  max_tokens: 777\n{alternates}");
695        let resolved = resolve(&manifest("m", &with_cap, control)).unwrap();
696        let caps: Vec<_> = resolved
697            .iter()
698            .map(|choice| (choice.choice.as_str(), choice.limits.max_output_tokens))
699            .collect();
700        assert_eq!(caps, [("primary", 777), ("cheap", 777), ("other", 777)]);
701        assert_eq!(resolved[1].limits.context_window, Some(2000));
702        assert!(resolved[1].warning.is_none());
703        assert!(resolved[2].warning.is_some());
704    }
705
706    #[test]
707    fn a_malformed_table_is_refused_naming_the_driver_and_field() {
708        let manifest = manifest("m", "", "");
709        let refused = resolve_model_limits(manifest.inference.as_ref().unwrap(), None, |_| {
710            DriverManifest {
711                label: "local/drv@0.1.0",
712                manifest_yaml: "name: drv\nversion: 0.1.0\nruntime: driver\nmodels:\n  m:\n    \
713                                max_output_tokens: 0\n    context_window: 1\n",
714            }
715        })
716        .unwrap_err();
717        let message = refused.to_string();
718        assert!(
719            message.starts_with("driver local/drv@0.1.0 ships a malformed models: table:"),
720            "{message}"
721        );
722        assert!(message.contains("models.m.max_output_tokens"), "{message}");
723    }
724
725    #[test]
726    fn a_hook_completion_caps_at_the_table_entry_else_the_primary_cap() {
727        let table = murmur_artifact::parse_model_limits(DRIVER).unwrap();
728        let lifecycle = HookInferencePurpose::Lifecycle;
729        let compaction = HookInferencePurpose::Compaction;
730        let caps = HookOutputCaps::new(table.clone(), 4096, None, None);
731        for purpose in [lifecycle, compaction] {
732            assert_eq!(caps.for_call("m", purpose), 65536);
733            assert_eq!(caps.for_call("alt-m", purpose), 1000);
734            assert_eq!(caps.for_call("unknown", purpose), 4096);
735        }
736        let capped = HookOutputCaps::new(table, 4096, Some(2000), None);
737        assert_eq!(capped.for_call("m", lifecycle), 65536);
738        assert_eq!(capped.for_call("unknown", lifecycle), 4096);
739        assert_eq!(capped.for_call("m", compaction), 2000);
740        assert_eq!(capped.for_call("unknown", compaction), 2000);
741    }
742
743    #[test]
744    fn a_hook_completions_budget_is_cut_from_its_own_cap() {
745        let table = murmur_artifact::parse_model_limits(WORKED_DRIVER).unwrap();
746        let caps = HookOutputCaps::new(table.clone(), 4096, Some(2000), None);
747        // min(922000, 1050000 − 2000) for a compaction call; the table's cap otherwise.
748        assert_eq!(caps.input_budget("ceiled", 2000), Some(922_000));
749        assert_eq!(caps.input_budget("ceiled", 128_000), Some(922_000));
750        assert_eq!(caps.input_budget("plain", 128_000), Some(922_000));
751        assert_eq!(caps.input_budget("plain", 6000), Some(1_044_000));
752        assert_eq!(caps.input_budget("unlisted", 4096), None, "no window");
753        let windowed = HookOutputCaps::new(table, 4096, None, Some(200_000));
754        assert_eq!(windowed.input_budget("ceiled", 8192), Some(191_808));
755        assert_eq!(windowed.input_budget("unlisted", 4096), Some(195_904));
756        assert_eq!(
757            windowed.input_budget("unlisted", 300_000),
758            Some(0),
759            "a model the launch could not check saturates"
760        );
761    }
762
763    /// A driver with the worked-numbers entries: `ceiled` 1050000 / 128000 / 922000, `plain`
764    /// 1050000 / 128000 / —, `tight` 1050000 / 128000 / 900000, `small` 1000000 / 8192 / —,
765    /// `alt-m` 2000 / 1000 / — and `full` 4000 / 4000 / —.
766    const WORKED_DRIVER: &str =
767        "name: drv\nversion: 0.1.0\nruntime: driver\nmodels:\n  ceiled:\n    \
768        max_output_tokens: 128000\n    context_window: 1050000\n    max_input_tokens: 922000\n  \
769        plain:\n    max_output_tokens: 128000\n    context_window: 1050000\n  tight:\n    \
770        max_output_tokens: 128000\n    context_window: 1050000\n    max_input_tokens: 900000\n  \
771        small:\n    max_output_tokens: 8192\n    context_window: 1000000\n  alt-m:\n    \
772        max_output_tokens: 1000\n    context_window: 2000\n  full:\n    \
773        max_output_tokens: 4000\n    context_window: 4000\n";
774
775    fn resolve_worked(
776        manifest: &RuntimeManifest,
777    ) -> Result<Vec<ChoiceModelLimits>, ModelLimitsError> {
778        resolve_model_limits(
779            manifest.inference.as_ref().unwrap(),
780            manifest.context.as_ref(),
781            |_| DriverManifest {
782                label: "local/drv@0.1.0",
783                manifest_yaml: WORKED_DRIVER,
784            },
785        )
786    }
787
788    #[test]
789    fn every_worked_numbers_row_resolves_its_budget_and_basis() {
790        for (model, extra, context, budget, basis, fires_at) in [
791            (
792                "ceiled",
793                "",
794                "",
795                922_000,
796                BudgetBasis::MaxInputTokens,
797                461_000,
798            ),
799            (
800                "plain",
801                "  max_tokens: 6000\n",
802                "",
803                1_044_000,
804                BudgetBasis::ContextWindowLessOutputCap,
805                522_000,
806            ),
807            (
808                "tight",
809                "",
810                "",
811                900_000,
812                BudgetBasis::MaxInputTokens,
813                450_000,
814            ),
815            (
816                "ceiled",
817                "  max_tokens: 8192\n",
818                "context:\n  max_tokens: 200000\n",
819                191_808,
820                BudgetBasis::ContextWindowLessOutputCap,
821                95_904,
822            ),
823        ] {
824            let resolved = resolve_worked(&manifest(model, extra, context)).unwrap();
825            let limits = resolved[0].limits;
826            assert_eq!(
827                (limits.input_budget, limits.input_budget_basis),
828                (Some(budget), basis),
829                "{model} {extra:?} {context:?}"
830            );
831            // Threshold 0.5 fires at half the budget, and not one token below it.
832            let occupancy = |tokens| crate::agent::SessionOccupancy::estimate(tokens);
833            assert!(occupancy(fires_at).crosses(budget, 0.5));
834            assert!(!occupancy(fires_at - 1).crosses(budget, 0.5));
835        }
836        let ceiled = resolve_worked(&manifest("ceiled", "", "")).unwrap()[0].limits;
837        assert_eq!(
838            ceiled,
839            ResolvedModelLimits {
840                max_output_tokens: 128_000,
841                max_output_tokens_source: LimitSource::Driver,
842                context_window: Some(1_050_000),
843                context_window_source: LimitSource::Driver,
844                max_input_tokens: Some(922_000),
845                max_input_tokens_source: LimitSource::Driver,
846                input_budget: Some(922_000),
847                input_budget_basis: BudgetBasis::MaxInputTokens,
848            }
849        );
850        let plain =
851            resolve_worked(&manifest("plain", "  max_tokens: 6000\n", "")).unwrap()[0].limits;
852        assert_eq!(
853            (plain.max_input_tokens, plain.max_input_tokens_source),
854            (None, LimitSource::None)
855        );
856        assert_eq!(
857            resolve_worked(&manifest(
858                "small",
859                "  max_tokens: 200000\n",
860                "context:\n  max_tokens: 200000\n"
861            )),
862            Err(ModelLimitsError::NoInputBudget {
863                choice: "primary".to_string(),
864                model: "small".to_string(),
865                driver: "local/drv@0.1.0".to_string(),
866                window: 200_000,
867                window_source: LimitSource::Manifest,
868                cap: 200_000,
869                cap_source: LimitSource::Manifest,
870                ceiling: None,
871            })
872        );
873    }
874
875    #[test]
876    fn the_input_ceiling_wins_a_tie_with_the_window_less_the_cap() {
877        let limits = resolve_worked(&manifest("ceiled", "", "")).unwrap()[0].limits;
878        assert_eq!(
879            limits.input_budget_for(128_000),
880            Some(InputBudget {
881                tokens: 922_000,
882                basis: BudgetBasis::MaxInputTokens
883            })
884        );
885        assert_eq!(
886            limits.input_budget_for(128_001),
887            Some(InputBudget {
888                tokens: 921_999,
889                basis: BudgetBasis::ContextWindowLessOutputCap
890            })
891        );
892        assert_eq!(
893            limits.input_budget_for(2000),
894            Some(InputBudget {
895                tokens: 922_000,
896                basis: BudgetBasis::MaxInputTokens
897            })
898        );
899        let windowless = ResolvedModelLimits {
900            context_window: None,
901            ..limits
902        };
903        assert_eq!(windowless.input_budget_for(1), None);
904    }
905
906    #[test]
907    fn a_cap_that_fills_the_window_is_refused_naming_every_number() {
908        for cap in [200_000, 300_000] {
909            let refused = resolve_worked(&manifest(
910                "plain",
911                &format!("  max_tokens: {cap}\n"),
912                "context:\n  max_tokens: 200000\n",
913            ))
914            .unwrap_err();
915            assert_eq!(
916                refused.to_string(),
917                format!(
918                    "driver choice 'primary' has no input budget: its context window of 200000 \
919                     (manifest) less its output cap of {cap} (manifest) leaves no input tokens; \
920                     input ceiling none (model 'plain', driver local/drv@0.1.0)"
921                )
922            );
923        }
924        let ceiled = resolve_worked(&manifest(
925            "ceiled",
926            "  max_tokens: 200000\n",
927            "context:\n  max_tokens: 200000\n",
928        ))
929        .unwrap_err()
930        .to_string();
931        assert!(
932            ceiled.ends_with("input ceiling 922000 (model 'ceiled', driver local/drv@0.1.0)"),
933            "{ceiled}"
934        );
935        let one = resolve_worked(&manifest(
936            "plain",
937            "  max_tokens: 199999\n",
938            "context:\n  max_tokens: 200000\n",
939        ))
940        .unwrap();
941        assert_eq!(one[0].limits.input_budget, Some(1));
942    }
943
944    #[test]
945    fn an_alternate_is_refused_on_its_own_window() {
946        let alternates =
947            "  max_tokens: 2000\n  alternates:\n    - name: cheap\n      driver:\n        \
948                          artifact: local/drv\n      model: alt-m\n";
949        let control = "control:\n  agent_settings: [inference.driver]\n";
950        let refused = resolve_worked(&manifest("ceiled", alternates, control)).unwrap_err();
951        assert_eq!(
952            refused,
953            ModelLimitsError::NoInputBudget {
954                choice: "cheap".to_string(),
955                model: "alt-m".to_string(),
956                driver: "local/drv@0.1.0".to_string(),
957                window: 2000,
958                window_source: LimitSource::Driver,
959                cap: 2000,
960                cap_source: LimitSource::Manifest,
961                ceiling: None,
962            }
963        );
964    }
965
966    #[test]
967    fn a_compaction_cap_that_fills_the_compaction_window_is_refused() {
968        // Inside the output ceiling (128000) but at the window `context.max_tokens` sets.
969        let refused = resolve_worked(&manifest(
970            "ceiled",
971            &format!("  max_tokens: 8192\n{}", compaction_cap(100_000, "")),
972            "context:\n  max_tokens: 100000\n",
973        ))
974        .unwrap_err();
975        assert_eq!(
976            refused,
977            ModelLimitsError::CompactionNoInputBudget {
978                model: "ceiled".to_string(),
979                driver: "local/drv@0.1.0".to_string(),
980                window: 100_000,
981                window_source: LimitSource::Manifest,
982                cap: 100_000,
983                ceiling: Some(922_000),
984            }
985        );
986        assert_eq!(
987            refused.to_string(),
988            "compaction calls have no input budget: the context window of 100000 (manifest) for \
989             compaction model 'ceiled' less inference.compaction.max_tokens of 100000 leaves no \
990             input tokens; input ceiling 922000 (driver local/drv@0.1.0)"
991        );
992        // The table's window for the compaction model, when the manifest sets none.
993        assert_eq!(
994            resolve_worked(&manifest(
995                "ceiled",
996                &compaction_cap(4000, "    model: full\n"),
997                ""
998            )),
999            Err(ModelLimitsError::CompactionNoInputBudget {
1000                model: "full".to_string(),
1001                driver: "local/drv@0.1.0".to_string(),
1002                window: 4000,
1003                window_source: LimitSource::Driver,
1004                cap: 4000,
1005                ceiling: None,
1006            })
1007        );
1008        assert!(resolve_worked(&manifest(
1009            "ceiled",
1010            &compaction_cap(1000, "    model: alt-m\n"),
1011            ""
1012        ))
1013        .is_ok());
1014        assert!(resolve_worked(&manifest(
1015            "ceiled",
1016            &format!("  max_tokens: 8192\n{}", compaction_cap(99_999, "")),
1017            "context:\n  max_tokens: 100000\n",
1018        ))
1019        .is_ok());
1020        // An unlisted compaction model with no context.max_tokens has no window to check.
1021        assert!(resolve_worked(&manifest(
1022            "ceiled",
1023            &compaction_cap(u32::MAX, "    model: unlisted\n"),
1024            ""
1025        ))
1026        .is_ok());
1027    }
1028
1029    fn compaction_cap(max_tokens: u32, model_line: &str) -> String {
1030        format!("  compaction:\n{model_line}    max_tokens: {max_tokens}\n")
1031    }
1032
1033    #[test]
1034    fn a_compaction_cap_above_the_ceiling_is_refused_naming_it() {
1035        let refused = resolve(&manifest("m", &compaction_cap(65537, ""), ""));
1036        assert_eq!(
1037            refused,
1038            Err(ModelLimitsError::CompactionCapAboveCeiling {
1039                max_tokens: 65537,
1040                ceiling: 65536,
1041                model: "m".to_string(),
1042                driver: "local/drv@0.1.0".to_string(),
1043            })
1044        );
1045        assert_eq!(
1046            refused.unwrap_err().to_string(),
1047            "inference.compaction.max_tokens (65537) is above the output ceiling of 65536 that \
1048             driver local/drv@0.1.0 declares for compaction model 'm'"
1049        );
1050        let at_ceiling = resolve(&manifest("m", &compaction_cap(65536, ""), "")).unwrap();
1051        assert_eq!(
1052            at_ceiling[0].limits.max_output_tokens, 65536,
1053            "the compaction cap leaves the agent-turn cap alone"
1054        );
1055    }
1056
1057    #[test]
1058    fn the_compaction_model_is_the_one_checked() {
1059        let model_line = |model: &str| format!("    model: {model}\n");
1060        // alt-m's ceiling is 1000, so a value fine for the primary model is refused for it.
1061        assert_eq!(
1062            resolve(&manifest(
1063                "m",
1064                &compaction_cap(2000, &model_line("alt-m")),
1065                ""
1066            )),
1067            Err(ModelLimitsError::CompactionCapAboveCeiling {
1068                max_tokens: 2000,
1069                ceiling: 1000,
1070                model: "alt-m".to_string(),
1071                driver: "local/drv@0.1.0".to_string(),
1072            })
1073        );
1074        // ...and a value above the primary model's ceiling passes when compaction uses another.
1075        assert!(resolve(&manifest(
1076            "alt-m",
1077            &compaction_cap(60000, &model_line("m")),
1078            ""
1079        ))
1080        .is_ok());
1081        // An unlisted compaction model has no ceiling to check against.
1082        assert!(resolve(&manifest(
1083            "m",
1084            &compaction_cap(u32::MAX, &model_line("unlisted")),
1085            ""
1086        ))
1087        .is_ok());
1088        assert!(resolve(&manifest(
1089            "M",
1090            &format!("  max_tokens: 4096\n{}", compaction_cap(u32::MAX, "")),
1091            ""
1092        ))
1093        .is_ok());
1094    }
1095
1096    #[test]
1097    fn the_source_wire_names() {
1098        assert_eq!(
1099            serde_json::to_value(ResolvedModelLimits {
1100                max_output_tokens: 1,
1101                max_output_tokens_source: LimitSource::Manifest,
1102                context_window: None,
1103                context_window_source: LimitSource::None,
1104                max_input_tokens: None,
1105                max_input_tokens_source: LimitSource::None,
1106                input_budget: None,
1107                input_budget_basis: BudgetBasis::None,
1108            })
1109            .unwrap(),
1110            serde_json::json!({
1111                "max_output_tokens": 1,
1112                "max_output_tokens_source": "manifest",
1113                "context_window": null,
1114                "context_window_source": "none",
1115                "max_input_tokens": null,
1116                "max_input_tokens_source": "none",
1117                "input_budget": null,
1118                "input_budget_basis": "none",
1119            })
1120        );
1121        assert_eq!(LimitSource::Driver.wire_name(), "driver");
1122        let text =
1123            serde_json::to_string(&resolve_worked(&manifest("ceiled", "", "")).unwrap()[0].limits)
1124                .unwrap();
1125        assert_eq!(
1126            text,
1127            "{\"max_output_tokens\":128000,\"max_output_tokens_source\":\"driver\",\
1128             \"context_window\":1050000,\"context_window_source\":\"driver\",\
1129             \"max_input_tokens\":922000,\"max_input_tokens_source\":\"driver\",\
1130             \"input_budget\":922000,\"input_budget_basis\":\"max_input_tokens\"}"
1131        );
1132        for (basis, name) in [
1133            (BudgetBasis::MaxInputTokens, "max_input_tokens"),
1134            (
1135                BudgetBasis::ContextWindowLessOutputCap,
1136                "context_window_less_output_cap",
1137            ),
1138            (BudgetBasis::ContextWindow, "context_window"),
1139            (BudgetBasis::None, "none"),
1140        ] {
1141            assert_eq!(basis.wire_name(), name);
1142            assert_eq!(
1143                serde_json::to_value(basis).unwrap(),
1144                serde_json::json!(name)
1145            );
1146        }
1147    }
1148}