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 pub vision: Option<bool>,
181}
182
183#[derive(Debug, Clone, Copy)]
185pub struct DriverManifest<'a> {
186 pub label: &'a str,
188 pub manifest_yaml: &'a str,
190}
191
192#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
196pub enum ModelLimitsError {
197 #[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 #[error("driver {driver} ships a malformed models: table: {reason}")]
210 MalformedTable { driver: String, reason: String },
211 #[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 #[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 #[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
259pub 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 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
383fn 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
419fn 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#[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#[derive(Debug, Clone, PartialEq, Eq, Default)]
486pub(crate) struct HookOutputCaps {
487 table: BTreeMap<String, ModelLimits>,
489 fallback: u32,
491 compaction: Option<u32>,
493 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 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 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 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 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 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 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 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 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 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 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 assert!(resolve(&manifest(
1076 "alt-m",
1077 &compaction_cap(60000, &model_line("m")),
1078 ""
1079 ))
1080 .is_ok());
1081 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}