1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
46#[serde(rename_all = "lowercase")]
47pub enum LimitSource {
48 Manifest,
50 Driver,
52 None,
54}
55
56impl LimitSource {
57 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
76#[serde(rename_all = "snake_case")]
77pub enum BudgetBasis {
78 MaxInputTokens,
80 ContextWindowLessOutputCap,
82 ContextWindow,
84 None,
86}
87
88impl BudgetBasis {
89 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
103pub struct InputBudget {
104 pub tokens: u32,
105 pub basis: BudgetBasis,
106}
107
108impl InputBudget {
109 #[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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
136pub struct ResolvedModelLimits {
137 pub max_output_tokens: u32,
140 pub max_output_tokens_source: LimitSource,
142 pub context_window: Option<u32>,
145 pub context_window_source: LimitSource,
146 pub max_input_tokens: Option<u32>,
148 pub max_input_tokens_source: LimitSource,
150 pub input_budget: Option<u32>,
153 pub input_budget_basis: BudgetBasis,
154}
155
156impl ResolvedModelLimits {
157 #[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#[derive(Debug, Clone, PartialEq, Eq)]
168pub struct ChoiceModelLimits {
169 pub choice: String,
171 pub model: String,
172 pub driver: String,
174 pub limits: ResolvedModelLimits,
175 pub warning: Option<String>,
178}
179
180#[derive(Debug, Clone, Copy)]
182pub struct DriverManifest<'a> {
183 pub label: &'a str,
185 pub manifest_yaml: &'a str,
187}
188
189#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
193pub enum ModelLimitsError {
194 #[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 #[error("driver {driver} ships a malformed models: table: {reason}")]
207 MalformedTable { driver: String, reason: String },
208 #[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 #[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 #[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
256pub 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 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
379fn 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
415fn 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#[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#[derive(Debug, Clone, PartialEq, Eq, Default)]
482pub(crate) struct HookOutputCaps {
483 table: BTreeMap<String, ModelLimits>,
485 fallback: u32,
487 compaction: Option<u32>,
489 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 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 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 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 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 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 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 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 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 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 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 assert!(resolve(&manifest(
1072 "alt-m",
1073 &compaction_cap(60000, &model_line("m")),
1074 ""
1075 ))
1076 .is_ok());
1077 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}