1use std::collections::HashMap;
8
9use leviath_core::{
10 Blueprint, ContextLayout, EvictionStrategy, Region, RegionKind, truncate_at_boundary,
11};
12
13use crate::ContextWindow;
14
15pub fn init_window_seeded(
23 window: &mut ContextWindow,
24 blueprint: &Blueprint,
25 seeds: &HashMap<String, String>,
26) {
27 for region_def in &blueprint.context_layout.regions {
28 let mut region = Region::new(
29 region_def.name.clone(),
30 region_def.kind.clone(),
31 region_def.max_tokens,
32 );
33 region.summarizable = region_def.summarizable;
34 region.admission = region_def.admission;
35 window.add_region(region);
36 }
37
38 if window.get_region("tool_results").is_none() {
39 let tool_region = Region::new("tool_results".to_string(), RegionKind::Temporary, 5000);
40 window.add_region(tool_region);
41 }
42
43 if window.get_region("conversation").is_none() {
44 let conv_region = Region::new(
45 "conversation".to_string(),
46 RegionKind::SlidingWindow {
47 max_items: 50,
48 eviction_strategy: EvictionStrategy::PerItem,
49 },
50 10000,
51 );
52 window.add_region(conv_region);
53 }
54
55 if window
60 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
61 .is_none()
62 {
63 window.add_region(Region::new(
64 crate::output_tool::FINAL_OUTPUT_REGION.to_string(),
65 RegionKind::Pinned,
66 crate::output_tool::FINAL_OUTPUT_REGION_TOKENS,
67 ));
68 }
69
70 for (name, content) in seeds {
71 let target = if name == "task" {
76 task_region_name(blueprint)
77 } else {
78 blueprint
79 .context_layout
80 .regions
81 .iter()
82 .find(|r| &r.name == name)
83 .map(|r| r.name.clone())
84 };
85 if let Some(region_name) = target {
86 let budget = window
91 .get_region(®ion_name)
92 .map(|r| r.max_tokens)
93 .unwrap_or(0);
94 let fitted = fit_seed_to_budget(content, budget);
95 let tokens = leviath_core::estimate_tokens(&fitted);
96 let _ = window.add_to_region(®ion_name, fitted, tokens);
97 }
98 }
99}
100
101const SEED_TRUNCATION_MARKER: &str =
103 "\n[...truncated by leviath: seed exceeded this region's budget]";
104
105fn fit_seed_to_budget(content: &str, max_tokens: usize) -> String {
109 let allowed = max_tokens.saturating_sub(1).saturating_mul(4);
112 if content.len() <= allowed {
113 return content.to_string();
114 }
115 let Some(room) = allowed.checked_sub(SEED_TRUNCATION_MARKER.len()) else {
118 return String::new();
119 };
120 format!(
121 "{}{SEED_TRUNCATION_MARKER}",
122 truncate_at_boundary(content, room)
123 )
124}
125
126fn task_region_name(blueprint: &Blueprint) -> Option<String> {
129 blueprint
130 .context_layout
131 .regions
132 .iter()
133 .find(|r| r.name == "task" && matches!(r.kind, RegionKind::Pinned))
134 .or_else(|| {
135 blueprint
136 .context_layout
137 .regions
138 .iter()
139 .find(|r| matches!(r.kind, RegionKind::Pinned))
140 })
141 .map(|r| r.name.clone())
142}
143
144pub fn init_window(window: &mut ContextWindow, blueprint: &Blueprint, task: &str) {
148 let seeds = HashMap::from([("task".to_string(), task.to_string())]);
149 init_window_seeded(window, blueprint, &seeds);
150}
151
152pub fn apply_layout(window: &mut ContextWindow, layout: &ContextLayout) {
157 let mut new_regions = Vec::new();
158 let mut kept: std::collections::HashSet<&str> = std::collections::HashSet::new();
159 for region_def in &layout.regions {
160 let mut new_region = Region::new(
161 region_def.name.clone(),
162 region_def.kind.clone(),
163 region_def.max_tokens,
164 );
165 new_region.summarizable = region_def.summarizable;
166 new_region.admission = region_def.admission;
167
168 if let Some(existing) = window.get_region(®ion_def.name) {
169 for entry in &existing.content {
175 let _ = new_region.carry_entry(entry.clone());
176 }
177 new_region.taint = existing.taint.clone();
180 }
181
182 kept.insert(region_def.name.as_str());
183 new_regions.push(new_region);
184 }
185
186 let always_visible = [
208 "conversation",
209 "tool_results",
210 crate::output_tool::FINAL_OUTPUT_REGION,
211 leviath_core::layout::STAGE_INSTRUCTIONS_REGION,
212 ];
213 let mut hidden = std::collections::HashSet::new();
214 for existing in &window.regions {
215 if kept.contains(existing.name.as_str()) {
216 continue;
217 }
218 let mut carried = Region::new(
219 existing.name.clone(),
220 existing.kind.clone(),
221 existing.max_tokens,
222 );
223 carried.summarizable = existing.summarizable;
224 carried.admission = existing.admission;
225 for entry in &existing.content {
228 let _ = carried.carry_entry(entry.clone());
229 }
230 carried.taint = existing.taint.clone();
231 if !always_visible.contains(&existing.name.as_str()) {
232 hidden.insert(existing.name.clone());
233 }
234 new_regions.push(carried);
235 }
236 window.hidden = hidden;
238
239 window.regions = new_regions;
240 window.current_tokens = window.calculate_tokens();
241}
242
243pub fn ensure_stage_instructions_region(window: &mut ContextWindow, prompts: &[Option<String>]) {
277 let declared = leviath_core::layout::STAGE_INSTRUCTIONS_REGION;
278 if window.get_region(declared).is_some() {
279 return;
280 }
281 let widest = prompts
283 .iter()
284 .flatten()
285 .map(|p| leviath_core::estimate_tokens(&format!("[Stage instructions: {p}]")))
286 .max();
287 let Some(widest) = widest.filter(|t| *t > 0) else {
288 return;
289 };
290 let ceiling = window.max_tokens / INSTRUCTIONS_SHARE_OF_WINDOW;
291 window.add_region(Region::new(
292 declared.to_string(),
293 RegionKind::Pinned,
294 widest.min(ceiling),
295 ));
296}
297
298const INSTRUCTIONS_SHARE_OF_WINDOW: usize = 4;
306
307#[cfg(test)]
308mod tests {
309 use super::{
310 SEED_TRUNCATION_MARKER, apply_layout, fit_seed_to_budget, init_window, init_window_seeded,
311 };
312 use crate::ContextWindow;
313 use leviath_core::{
314 Blueprint, ContextLayout, EvictionStrategy, RegionKind, Stage, blueprint::ModelConfig,
315 layout::RegionDefinition,
316 };
317 use std::collections::HashMap;
318
319 fn blueprint_with(regions: Vec<RegionDefinition>) -> Blueprint {
320 let layout = ContextLayout::new(regions, 100_000);
321 let stages = vec![Stage::new(
322 "main".to_string(),
323 ModelConfig::new("anthropic".to_string(), "claude-sonnet-4".to_string()),
324 )];
325 Blueprint::new("bp".to_string(), "desc".to_string(), stages, layout)
326 }
327
328 fn seeded_window(bp: &Blueprint, task: &str) -> ContextWindow {
329 let mut window = ContextWindow::new(100_000);
330 init_window(&mut window, bp, task);
331 window
332 }
333
334 #[test]
338 fn a_layout_that_declares_final_output_keeps_its_own() {
339 const DECLARED_TOKENS: usize = 12_345;
340 let bp = blueprint_with(vec![
341 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 1_000),
342 RegionDefinition::new(
343 crate::output_tool::FINAL_OUTPUT_REGION.to_string(),
344 RegionKind::Pinned,
345 DECLARED_TOKENS,
346 ),
347 ]);
348
349 let window = seeded_window(&bp, "t");
350
351 assert_eq!(
352 window
353 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
354 .expect("the region is there")
355 .max_tokens,
356 DECLARED_TOKENS,
357 "the blueprint's own budget survives"
358 );
359 }
360
361 #[test]
364 fn a_layout_without_final_output_gets_the_default_one() {
365 let bp = blueprint_with(vec![RegionDefinition::new(
366 "task".to_string(),
367 RegionKind::Pinned,
368 1_000,
369 )]);
370
371 let window = seeded_window(&bp, "t");
372
373 assert_eq!(
374 window
375 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
376 .expect("added for us")
377 .max_tokens,
378 crate::output_tool::FINAL_OUTPUT_REGION_TOKENS
379 );
380 }
381
382 #[test]
383 fn init_window_seeded_fills_multiple_named_regions_and_ignores_unknown() {
384 let bp = blueprint_with(vec![
385 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
386 RegionDefinition::new("criteria".to_string(), RegionKind::Pinned, 5000),
387 ]);
388 let seeds = HashMap::from([
389 ("task".to_string(), "build a parser".to_string()),
390 ("criteria".to_string(), "focus on safety".to_string()),
391 ("ghost".to_string(), "no such region".to_string()),
392 ]);
393 let mut window = ContextWindow::new(100_000);
394 init_window_seeded(&mut window, &bp, &seeds);
395
396 assert!(
397 window
398 .get_region("task")
399 .unwrap()
400 .content
401 .iter()
402 .any(|e| e.content.contains("build a parser"))
403 );
404 assert!(
405 window
406 .get_region("criteria")
407 .unwrap()
408 .content
409 .iter()
410 .any(|e| e.content.contains("focus on safety"))
411 );
412 assert!(window.get_region("ghost").is_none());
414 }
415
416 #[test]
417 fn fit_seed_to_budget_leaves_a_fitting_seed_untouched() {
418 assert_eq!(fit_seed_to_budget("hello", 100), "hello");
419 let exact = "x".repeat(36);
421 assert_eq!(fit_seed_to_budget(&exact, 10), exact);
422 }
423
424 fn estimated_tokens(fitted: &str) -> usize {
427 leviath_core::estimate_tokens(fitted)
428 }
429
430 #[test]
431 fn fit_seed_to_budget_truncates_and_marks_an_oversized_seed() {
432 let big = "x".repeat(10_000);
433 let fitted = fit_seed_to_budget(&big, 100);
434 assert!(fitted.ends_with(SEED_TRUNCATION_MARKER));
435 let estimate = estimated_tokens(&fitted);
437 assert!(estimate <= 100, "estimate was {estimate}");
438 }
439
440 #[test]
441 fn fit_seed_to_budget_cuts_on_a_char_boundary() {
442 const MAX_TOKENS: usize = 60;
445 let room = (MAX_TOKENS - 1) * 4 - SEED_TRUNCATION_MARKER.len();
446 let mut s = "a".repeat(room - 1);
447 s.push('é'); s.push_str(&"b".repeat(500));
449 assert!(!s.is_char_boundary(room), "test must straddle the cut");
450
451 let fitted = fit_seed_to_budget(&s, MAX_TOKENS);
452 assert!(fitted.ends_with(SEED_TRUNCATION_MARKER));
453 assert!(estimated_tokens(&fitted) <= MAX_TOKENS);
454 assert_eq!(
456 fitted,
457 format!("{}{SEED_TRUNCATION_MARKER}", "a".repeat(room - 1))
458 );
459 }
460
461 #[test]
462 fn fit_seed_to_budget_yields_nothing_when_even_the_marker_cannot_fit() {
463 assert_eq!(fit_seed_to_budget("some content here", 2), "");
466 assert_eq!(fit_seed_to_budget("x", 0), "");
468 }
469
470 #[test]
471 fn init_window_seeded_truncates_a_seed_larger_than_its_region() {
472 let bp = blueprint_with(vec![RegionDefinition::new(
476 "facts".to_string(),
477 RegionKind::Pinned,
478 50,
479 )]);
480 let seeds = HashMap::from([("facts".to_string(), "y".repeat(10_000))]);
481 let mut window = ContextWindow::new(100_000);
482 init_window_seeded(&mut window, &bp, &seeds);
483
484 let region = window.get_region("facts").unwrap();
485 assert!(
486 !region.content.is_empty(),
487 "an oversized seed must be trimmed, not dropped"
488 );
489 assert!(region.content[0].content.ends_with(SEED_TRUNCATION_MARKER));
490 }
491
492 #[test]
493 fn init_window_seeded_task_key_falls_back_to_first_pinned() {
494 let bp = blueprint_with(vec![RegionDefinition::new(
497 "system".to_string(),
498 RegionKind::Pinned,
499 5000,
500 )]);
501 let seeds = HashMap::from([("task".to_string(), "fallback text".to_string())]);
502 let mut window = ContextWindow::new(100_000);
503 init_window_seeded(&mut window, &bp, &seeds);
504 assert!(
505 window
506 .get_region("system")
507 .unwrap()
508 .content
509 .iter()
510 .any(|e| e.content.contains("fallback text"))
511 );
512 }
513
514 #[test]
515 fn init_prefers_named_task_region_and_keeps_existing_infra_regions() {
516 let bp = blueprint_with(vec![
517 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
518 RegionDefinition::new("tool_results".to_string(), RegionKind::Temporary, 5000),
519 RegionDefinition::new(
520 "conversation".to_string(),
521 RegionKind::SlidingWindow {
522 max_items: 10,
523 eviction_strategy: EvictionStrategy::PerItem,
524 },
525 10_000,
526 ),
527 ]);
528
529 let window = seeded_window(&bp, "do the thing");
530 assert!(
532 window
533 .get_region("task")
534 .unwrap()
535 .content
536 .iter()
537 .any(|e| e.content.contains("do the thing"))
538 );
539 assert_eq!(
541 window
542 .regions
543 .iter()
544 .filter(|r| r.name == "tool_results")
545 .count(),
546 1
547 );
548 assert_eq!(
549 window
550 .regions
551 .iter()
552 .filter(|r| r.name == "conversation")
553 .count(),
554 1
555 );
556 }
557
558 #[test]
559 fn init_adds_infra_regions_and_falls_back_to_first_pinned() {
560 let bp = blueprint_with(vec![RegionDefinition::new(
563 "system".to_string(),
564 RegionKind::Pinned,
565 5000,
566 )]);
567
568 let window = seeded_window(&bp, "seed task");
569 assert!(window.get_region("tool_results").is_some());
570 assert!(window.get_region("conversation").is_some());
571 assert!(
572 window
573 .get_region("system")
574 .unwrap()
575 .content
576 .iter()
577 .any(|e| e.content.contains("seed task"))
578 );
579 }
580
581 #[test]
582 fn init_without_pinned_region_does_not_seed_task() {
583 let bp = blueprint_with(vec![RegionDefinition::new(
584 "scratch".to_string(),
585 RegionKind::Temporary,
586 5000,
587 )]);
588
589 let window = seeded_window(&bp, "unseeded task");
590 assert!(window.get_region("scratch").unwrap().content.is_empty());
593 assert!(window.get_region("tool_results").is_some());
595 assert!(window.get_region("conversation").is_some());
596 }
597
598 #[test]
599 fn init_task_named_region_that_is_not_pinned_falls_back_to_first_pinned() {
600 let bp = blueprint_with(vec![
604 RegionDefinition::new("task".to_string(), RegionKind::Temporary, 5000),
605 RegionDefinition::new("system".to_string(), RegionKind::Pinned, 5000),
606 ]);
607
608 let window = seeded_window(&bp, "fallback seed");
609 assert!(window.get_region("task").unwrap().content.is_empty());
611 assert!(
613 window
614 .get_region("system")
615 .unwrap()
616 .content
617 .iter()
618 .any(|e| e.content.contains("fallback seed"))
619 );
620 }
621
622 #[test]
623 fn apply_layout_preserves_overlapping_content_and_creates_new_regions() {
624 let bp = blueprint_with(vec![RegionDefinition::new(
625 "system".to_string(),
626 RegionKind::Pinned,
627 5000,
628 )]);
629 let mut window = seeded_window(&bp, "carried content");
630
631 let new_layout = ContextLayout::new(
634 vec![
635 RegionDefinition::new("system".to_string(), RegionKind::Pinned, 5000),
636 RegionDefinition::new("scratch".to_string(), RegionKind::Temporary, 3000),
637 ],
638 8000,
639 );
640
641 apply_layout(&mut window, &new_layout);
642
643 assert_eq!(window.regions.len(), 5);
649 assert!(window.get_region("conversation").is_some());
650 assert!(window.get_region("tool_results").is_some());
651 assert!(
652 window
653 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
654 .is_some(),
655 "a submitted answer must survive a stage transition"
656 );
657 assert!(
658 window
659 .get_region("system")
660 .unwrap()
661 .content
662 .iter()
663 .any(|e| e.content.contains("carried content"))
664 );
665 assert!(window.get_region("scratch").unwrap().content.is_empty());
666 assert_eq!(window.current_tokens, window.calculate_tokens());
668 assert!(window.current_tokens > 0);
669 }
670
671 #[test]
672 fn apply_layout_preserves_entry_kinds_and_taint_across_swap() {
673 let bp = blueprint_with(vec![RegionDefinition::new(
677 "task".to_string(),
678 RegionKind::Pinned,
679 5000,
680 )]);
681 let mut window = seeded_window(&bp, "the task");
682 window
683 .add_typed_entry(
684 "conversation",
685 leviath_core::EntryKind::AssistantTurn {
686 tool_calls: vec![leviath_core::SerializedToolCall {
687 id: "call_9".to_string(),
688 name: "shell".to_string(),
689 arguments: serde_json::json!({"command": "ls"}),
690 thought_signature: None,
691 }],
692 },
693 "running ls".to_string(),
694 10,
695 )
696 .unwrap();
697 window
698 .add_typed_entry(
699 "conversation",
700 leviath_core::EntryKind::ToolResult {
701 tool_call_id: "call_9".to_string(),
702 tool_name: "shell".to_string(),
703 is_error: false,
704 },
705 "file_a\nfile_b".to_string(),
706 10,
707 )
708 .unwrap();
709 window
710 .get_region_mut("conversation")
711 .unwrap()
712 .enable_taint_tracking();
713
714 let omitting = ContextLayout::new(
716 vec![RegionDefinition::new(
717 "task".to_string(),
718 RegionKind::Pinned,
719 5000,
720 )],
721 8000,
722 );
723 apply_layout(&mut window, &omitting);
724
725 let declaring = ContextLayout::new(
727 vec![
728 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
729 RegionDefinition::new(
730 "conversation".to_string(),
731 RegionKind::SlidingWindow {
732 max_items: 10,
733 eviction_strategy: EvictionStrategy::PerItem,
734 },
735 10_000,
736 ),
737 ],
738 20_000,
739 );
740 apply_layout(&mut window, &declaring);
741
742 let conv = window.get_region("conversation").unwrap();
743 assert!(
744 conv.content.iter().any(|e| matches!(
745 &e.kind,
746 leviath_core::EntryKind::AssistantTurn { tool_calls }
747 if tool_calls.iter().any(|c| c.id == "call_9")
748 )),
749 "assistant turn must keep its typed tool_calls through both carry paths"
750 );
751 assert!(
752 conv.content.iter().any(|e| matches!(
753 &e.kind,
754 leviath_core::EntryKind::ToolResult { tool_call_id, .. }
755 if tool_call_id == "call_9"
756 )),
757 "tool result must keep its typed pairing through both carry paths"
758 );
759 assert!(
760 conv.taint.is_some(),
761 "region-level taint state must carry across layout swaps"
762 );
763 }
764
765 #[test]
766 fn apply_layout_carries_conversation_when_new_layout_omits_it() {
767 let bp = blueprint_with(vec![RegionDefinition::new(
771 "task".to_string(),
772 RegionKind::Pinned,
773 5000,
774 )]);
775 let mut window = seeded_window(&bp, "the task");
776 window
777 .add_typed_entry(
778 "conversation",
779 leviath_core::EntryKind::UserMessage,
780 "hello from stage 0".to_string(),
781 10,
782 )
783 .unwrap();
784
785 let next = ContextLayout::new(
787 vec![RegionDefinition::new(
788 "task".to_string(),
789 RegionKind::Pinned,
790 5000,
791 )],
792 8000,
793 );
794 apply_layout(&mut window, &next);
795
796 let conv = window
797 .get_region("conversation")
798 .expect("conversation carried across transition");
799 assert!(
800 conv.content
801 .iter()
802 .any(|e| e.content.contains("hello from stage 0")),
803 "carried conversation must retain its history"
804 );
805 }
806}