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 window.add_region(region);
35 }
36
37 if window.get_region("tool_results").is_none() {
38 let tool_region = Region::new("tool_results".to_string(), RegionKind::Temporary, 5000);
39 window.add_region(tool_region);
40 }
41
42 if window.get_region("conversation").is_none() {
43 let conv_region = Region::new(
44 "conversation".to_string(),
45 RegionKind::SlidingWindow {
46 max_items: 50,
47 eviction_strategy: EvictionStrategy::PerItem,
48 },
49 10000,
50 );
51 window.add_region(conv_region);
52 }
53
54 if window
59 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
60 .is_none()
61 {
62 window.add_region(Region::new(
63 crate::output_tool::FINAL_OUTPUT_REGION.to_string(),
64 RegionKind::Pinned,
65 crate::output_tool::FINAL_OUTPUT_REGION_TOKENS,
66 ));
67 }
68
69 for (name, content) in seeds {
70 let target = if name == "task" {
75 task_region_name(blueprint)
76 } else {
77 blueprint
78 .context_layout
79 .regions
80 .iter()
81 .find(|r| &r.name == name)
82 .map(|r| r.name.clone())
83 };
84 if let Some(region_name) = target {
85 let budget = window
90 .get_region(®ion_name)
91 .map(|r| r.max_tokens)
92 .unwrap_or(0);
93 let fitted = fit_seed_to_budget(content, budget);
94 let tokens = leviath_core::estimate_tokens(&fitted);
95 let _ = window.add_to_region(®ion_name, fitted, tokens);
96 }
97 }
98}
99
100const SEED_TRUNCATION_MARKER: &str =
102 "\n[...truncated by leviath: seed exceeded this region's budget]";
103
104fn fit_seed_to_budget(content: &str, max_tokens: usize) -> String {
108 let allowed = max_tokens.saturating_sub(1).saturating_mul(4);
111 if content.len() <= allowed {
112 return content.to_string();
113 }
114 let Some(room) = allowed.checked_sub(SEED_TRUNCATION_MARKER.len()) else {
117 return String::new();
118 };
119 format!(
120 "{}{SEED_TRUNCATION_MARKER}",
121 truncate_at_boundary(content, room)
122 )
123}
124
125fn task_region_name(blueprint: &Blueprint) -> Option<String> {
128 blueprint
129 .context_layout
130 .regions
131 .iter()
132 .find(|r| r.name == "task" && matches!(r.kind, RegionKind::Pinned))
133 .or_else(|| {
134 blueprint
135 .context_layout
136 .regions
137 .iter()
138 .find(|r| matches!(r.kind, RegionKind::Pinned))
139 })
140 .map(|r| r.name.clone())
141}
142
143pub fn init_window(window: &mut ContextWindow, blueprint: &Blueprint, task: &str) {
147 let seeds = HashMap::from([("task".to_string(), task.to_string())]);
148 init_window_seeded(window, blueprint, &seeds);
149}
150
151pub fn apply_layout(window: &mut ContextWindow, layout: &ContextLayout) {
156 let mut new_regions = Vec::new();
157 let mut kept: std::collections::HashSet<&str> = std::collections::HashSet::new();
158 for region_def in &layout.regions {
159 let mut new_region = Region::new(
160 region_def.name.clone(),
161 region_def.kind.clone(),
162 region_def.max_tokens,
163 );
164 new_region.summarizable = region_def.summarizable;
165
166 if let Some(existing) = window.get_region(®ion_def.name) {
167 for entry in &existing.content {
173 let _ = new_region.carry_entry(entry.clone());
174 }
175 new_region.taint = existing.taint.clone();
178 }
179
180 kept.insert(region_def.name.as_str());
181 new_regions.push(new_region);
182 }
183
184 let always_visible = [
206 "conversation",
207 "tool_results",
208 crate::output_tool::FINAL_OUTPUT_REGION,
209 leviath_core::layout::STAGE_INSTRUCTIONS_REGION,
210 ];
211 let mut hidden = std::collections::HashSet::new();
212 for existing in &window.regions {
213 if kept.contains(existing.name.as_str()) {
214 continue;
215 }
216 let mut carried = Region::new(
217 existing.name.clone(),
218 existing.kind.clone(),
219 existing.max_tokens,
220 );
221 carried.summarizable = existing.summarizable;
222 for entry in &existing.content {
225 let _ = carried.carry_entry(entry.clone());
226 }
227 carried.taint = existing.taint.clone();
228 if !always_visible.contains(&existing.name.as_str()) {
229 hidden.insert(existing.name.clone());
230 }
231 new_regions.push(carried);
232 }
233 window.hidden = hidden;
235
236 window.regions = new_regions;
237 window.current_tokens = window.calculate_tokens();
238}
239
240#[cfg(test)]
241mod tests {
242 use super::{
243 SEED_TRUNCATION_MARKER, apply_layout, fit_seed_to_budget, init_window, init_window_seeded,
244 };
245 use crate::ContextWindow;
246 use leviath_core::{
247 Blueprint, ContextLayout, EvictionStrategy, RegionKind, Stage, blueprint::ModelConfig,
248 layout::RegionDefinition,
249 };
250 use std::collections::HashMap;
251
252 fn blueprint_with(regions: Vec<RegionDefinition>) -> Blueprint {
253 let layout = ContextLayout::new(regions, 100_000);
254 let stages = vec![Stage::new(
255 "main".to_string(),
256 ModelConfig::new("anthropic".to_string(), "claude-sonnet-4".to_string()),
257 )];
258 Blueprint::new("bp".to_string(), "desc".to_string(), stages, layout)
259 }
260
261 fn seeded_window(bp: &Blueprint, task: &str) -> ContextWindow {
262 let mut window = ContextWindow::new(100_000);
263 init_window(&mut window, bp, task);
264 window
265 }
266
267 #[test]
271 fn a_layout_that_declares_final_output_keeps_its_own() {
272 const DECLARED_TOKENS: usize = 12_345;
273 let bp = blueprint_with(vec![
274 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 1_000),
275 RegionDefinition::new(
276 crate::output_tool::FINAL_OUTPUT_REGION.to_string(),
277 RegionKind::Pinned,
278 DECLARED_TOKENS,
279 ),
280 ]);
281
282 let window = seeded_window(&bp, "t");
283
284 assert_eq!(
285 window
286 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
287 .expect("the region is there")
288 .max_tokens,
289 DECLARED_TOKENS,
290 "the blueprint's own budget survives"
291 );
292 }
293
294 #[test]
297 fn a_layout_without_final_output_gets_the_default_one() {
298 let bp = blueprint_with(vec![RegionDefinition::new(
299 "task".to_string(),
300 RegionKind::Pinned,
301 1_000,
302 )]);
303
304 let window = seeded_window(&bp, "t");
305
306 assert_eq!(
307 window
308 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
309 .expect("added for us")
310 .max_tokens,
311 crate::output_tool::FINAL_OUTPUT_REGION_TOKENS
312 );
313 }
314
315 #[test]
316 fn init_window_seeded_fills_multiple_named_regions_and_ignores_unknown() {
317 let bp = blueprint_with(vec![
318 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
319 RegionDefinition::new("criteria".to_string(), RegionKind::Pinned, 5000),
320 ]);
321 let seeds = HashMap::from([
322 ("task".to_string(), "build a parser".to_string()),
323 ("criteria".to_string(), "focus on safety".to_string()),
324 ("ghost".to_string(), "no such region".to_string()),
325 ]);
326 let mut window = ContextWindow::new(100_000);
327 init_window_seeded(&mut window, &bp, &seeds);
328
329 assert!(
330 window
331 .get_region("task")
332 .unwrap()
333 .content
334 .iter()
335 .any(|e| e.content.contains("build a parser"))
336 );
337 assert!(
338 window
339 .get_region("criteria")
340 .unwrap()
341 .content
342 .iter()
343 .any(|e| e.content.contains("focus on safety"))
344 );
345 assert!(window.get_region("ghost").is_none());
347 }
348
349 #[test]
350 fn fit_seed_to_budget_leaves_a_fitting_seed_untouched() {
351 assert_eq!(fit_seed_to_budget("hello", 100), "hello");
352 let exact = "x".repeat(36);
354 assert_eq!(fit_seed_to_budget(&exact, 10), exact);
355 }
356
357 fn estimated_tokens(fitted: &str) -> usize {
360 leviath_core::estimate_tokens(fitted)
361 }
362
363 #[test]
364 fn fit_seed_to_budget_truncates_and_marks_an_oversized_seed() {
365 let big = "x".repeat(10_000);
366 let fitted = fit_seed_to_budget(&big, 100);
367 assert!(fitted.ends_with(SEED_TRUNCATION_MARKER));
368 let estimate = estimated_tokens(&fitted);
370 assert!(estimate <= 100, "estimate was {estimate}");
371 }
372
373 #[test]
374 fn fit_seed_to_budget_cuts_on_a_char_boundary() {
375 const MAX_TOKENS: usize = 60;
378 let room = (MAX_TOKENS - 1) * 4 - SEED_TRUNCATION_MARKER.len();
379 let mut s = "a".repeat(room - 1);
380 s.push('é'); s.push_str(&"b".repeat(500));
382 assert!(!s.is_char_boundary(room), "test must straddle the cut");
383
384 let fitted = fit_seed_to_budget(&s, MAX_TOKENS);
385 assert!(fitted.ends_with(SEED_TRUNCATION_MARKER));
386 assert!(estimated_tokens(&fitted) <= MAX_TOKENS);
387 assert_eq!(
389 fitted,
390 format!("{}{SEED_TRUNCATION_MARKER}", "a".repeat(room - 1))
391 );
392 }
393
394 #[test]
395 fn fit_seed_to_budget_yields_nothing_when_even_the_marker_cannot_fit() {
396 assert_eq!(fit_seed_to_budget("some content here", 2), "");
399 assert_eq!(fit_seed_to_budget("x", 0), "");
401 }
402
403 #[test]
404 fn init_window_seeded_truncates_a_seed_larger_than_its_region() {
405 let bp = blueprint_with(vec![RegionDefinition::new(
409 "facts".to_string(),
410 RegionKind::Pinned,
411 50,
412 )]);
413 let seeds = HashMap::from([("facts".to_string(), "y".repeat(10_000))]);
414 let mut window = ContextWindow::new(100_000);
415 init_window_seeded(&mut window, &bp, &seeds);
416
417 let region = window.get_region("facts").unwrap();
418 assert!(
419 !region.content.is_empty(),
420 "an oversized seed must be trimmed, not dropped"
421 );
422 assert!(region.content[0].content.ends_with(SEED_TRUNCATION_MARKER));
423 }
424
425 #[test]
426 fn init_window_seeded_task_key_falls_back_to_first_pinned() {
427 let bp = blueprint_with(vec![RegionDefinition::new(
430 "system".to_string(),
431 RegionKind::Pinned,
432 5000,
433 )]);
434 let seeds = HashMap::from([("task".to_string(), "fallback text".to_string())]);
435 let mut window = ContextWindow::new(100_000);
436 init_window_seeded(&mut window, &bp, &seeds);
437 assert!(
438 window
439 .get_region("system")
440 .unwrap()
441 .content
442 .iter()
443 .any(|e| e.content.contains("fallback text"))
444 );
445 }
446
447 #[test]
448 fn init_prefers_named_task_region_and_keeps_existing_infra_regions() {
449 let bp = blueprint_with(vec![
450 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
451 RegionDefinition::new("tool_results".to_string(), RegionKind::Temporary, 5000),
452 RegionDefinition::new(
453 "conversation".to_string(),
454 RegionKind::SlidingWindow {
455 max_items: 10,
456 eviction_strategy: EvictionStrategy::PerItem,
457 },
458 10_000,
459 ),
460 ]);
461
462 let window = seeded_window(&bp, "do the thing");
463 assert!(
465 window
466 .get_region("task")
467 .unwrap()
468 .content
469 .iter()
470 .any(|e| e.content.contains("do the thing"))
471 );
472 assert_eq!(
474 window
475 .regions
476 .iter()
477 .filter(|r| r.name == "tool_results")
478 .count(),
479 1
480 );
481 assert_eq!(
482 window
483 .regions
484 .iter()
485 .filter(|r| r.name == "conversation")
486 .count(),
487 1
488 );
489 }
490
491 #[test]
492 fn init_adds_infra_regions_and_falls_back_to_first_pinned() {
493 let bp = blueprint_with(vec![RegionDefinition::new(
496 "system".to_string(),
497 RegionKind::Pinned,
498 5000,
499 )]);
500
501 let window = seeded_window(&bp, "seed task");
502 assert!(window.get_region("tool_results").is_some());
503 assert!(window.get_region("conversation").is_some());
504 assert!(
505 window
506 .get_region("system")
507 .unwrap()
508 .content
509 .iter()
510 .any(|e| e.content.contains("seed task"))
511 );
512 }
513
514 #[test]
515 fn init_without_pinned_region_does_not_seed_task() {
516 let bp = blueprint_with(vec![RegionDefinition::new(
517 "scratch".to_string(),
518 RegionKind::Temporary,
519 5000,
520 )]);
521
522 let window = seeded_window(&bp, "unseeded task");
523 assert!(window.get_region("scratch").unwrap().content.is_empty());
526 assert!(window.get_region("tool_results").is_some());
528 assert!(window.get_region("conversation").is_some());
529 }
530
531 #[test]
532 fn init_task_named_region_that_is_not_pinned_falls_back_to_first_pinned() {
533 let bp = blueprint_with(vec![
537 RegionDefinition::new("task".to_string(), RegionKind::Temporary, 5000),
538 RegionDefinition::new("system".to_string(), RegionKind::Pinned, 5000),
539 ]);
540
541 let window = seeded_window(&bp, "fallback seed");
542 assert!(window.get_region("task").unwrap().content.is_empty());
544 assert!(
546 window
547 .get_region("system")
548 .unwrap()
549 .content
550 .iter()
551 .any(|e| e.content.contains("fallback seed"))
552 );
553 }
554
555 #[test]
556 fn apply_layout_preserves_overlapping_content_and_creates_new_regions() {
557 let bp = blueprint_with(vec![RegionDefinition::new(
558 "system".to_string(),
559 RegionKind::Pinned,
560 5000,
561 )]);
562 let mut window = seeded_window(&bp, "carried content");
563
564 let new_layout = ContextLayout::new(
567 vec![
568 RegionDefinition::new("system".to_string(), RegionKind::Pinned, 5000),
569 RegionDefinition::new("scratch".to_string(), RegionKind::Temporary, 3000),
570 ],
571 8000,
572 );
573
574 apply_layout(&mut window, &new_layout);
575
576 assert_eq!(window.regions.len(), 5);
582 assert!(window.get_region("conversation").is_some());
583 assert!(window.get_region("tool_results").is_some());
584 assert!(
585 window
586 .get_region(crate::output_tool::FINAL_OUTPUT_REGION)
587 .is_some(),
588 "a submitted answer must survive a stage transition"
589 );
590 assert!(
591 window
592 .get_region("system")
593 .unwrap()
594 .content
595 .iter()
596 .any(|e| e.content.contains("carried content"))
597 );
598 assert!(window.get_region("scratch").unwrap().content.is_empty());
599 assert_eq!(window.current_tokens, window.calculate_tokens());
601 assert!(window.current_tokens > 0);
602 }
603
604 #[test]
605 fn apply_layout_preserves_entry_kinds_and_taint_across_swap() {
606 let bp = blueprint_with(vec![RegionDefinition::new(
610 "task".to_string(),
611 RegionKind::Pinned,
612 5000,
613 )]);
614 let mut window = seeded_window(&bp, "the task");
615 window
616 .add_typed_entry(
617 "conversation",
618 leviath_core::EntryKind::AssistantTurn {
619 tool_calls: vec![leviath_core::SerializedToolCall {
620 id: "call_9".to_string(),
621 name: "shell".to_string(),
622 arguments: serde_json::json!({"command": "ls"}),
623 thought_signature: None,
624 }],
625 },
626 "running ls".to_string(),
627 10,
628 )
629 .unwrap();
630 window
631 .add_typed_entry(
632 "conversation",
633 leviath_core::EntryKind::ToolResult {
634 tool_call_id: "call_9".to_string(),
635 tool_name: "shell".to_string(),
636 is_error: false,
637 },
638 "file_a\nfile_b".to_string(),
639 10,
640 )
641 .unwrap();
642 window
643 .get_region_mut("conversation")
644 .unwrap()
645 .enable_taint_tracking();
646
647 let omitting = ContextLayout::new(
649 vec![RegionDefinition::new(
650 "task".to_string(),
651 RegionKind::Pinned,
652 5000,
653 )],
654 8000,
655 );
656 apply_layout(&mut window, &omitting);
657
658 let declaring = ContextLayout::new(
660 vec![
661 RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
662 RegionDefinition::new(
663 "conversation".to_string(),
664 RegionKind::SlidingWindow {
665 max_items: 10,
666 eviction_strategy: EvictionStrategy::PerItem,
667 },
668 10_000,
669 ),
670 ],
671 20_000,
672 );
673 apply_layout(&mut window, &declaring);
674
675 let conv = window.get_region("conversation").unwrap();
676 assert!(
677 conv.content.iter().any(|e| matches!(
678 &e.kind,
679 leviath_core::EntryKind::AssistantTurn { tool_calls }
680 if tool_calls.iter().any(|c| c.id == "call_9")
681 )),
682 "assistant turn must keep its typed tool_calls through both carry paths"
683 );
684 assert!(
685 conv.content.iter().any(|e| matches!(
686 &e.kind,
687 leviath_core::EntryKind::ToolResult { tool_call_id, .. }
688 if tool_call_id == "call_9"
689 )),
690 "tool result must keep its typed pairing through both carry paths"
691 );
692 assert!(
693 conv.taint.is_some(),
694 "region-level taint state must carry across layout swaps"
695 );
696 }
697
698 #[test]
699 fn apply_layout_carries_conversation_when_new_layout_omits_it() {
700 let bp = blueprint_with(vec![RegionDefinition::new(
704 "task".to_string(),
705 RegionKind::Pinned,
706 5000,
707 )]);
708 let mut window = seeded_window(&bp, "the task");
709 window
710 .add_typed_entry(
711 "conversation",
712 leviath_core::EntryKind::UserMessage,
713 "hello from stage 0".to_string(),
714 10,
715 )
716 .unwrap();
717
718 let next = ContextLayout::new(
720 vec![RegionDefinition::new(
721 "task".to_string(),
722 RegionKind::Pinned,
723 5000,
724 )],
725 8000,
726 );
727 apply_layout(&mut window, &next);
728
729 let conv = window
730 .get_region("conversation")
731 .expect("conversation carried across transition");
732 assert!(
733 conv.content
734 .iter()
735 .any(|e| e.content.contains("hello from stage 0")),
736 "carried conversation must retain its history"
737 );
738 }
739}