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