use std::collections::HashMap;
use leviath_core::{
Blueprint, ContextLayout, EvictionStrategy, Region, RegionKind, truncate_at_boundary,
};
use crate::ContextWindow;
pub fn init_window_seeded(
window: &mut ContextWindow,
blueprint: &Blueprint,
seeds: &HashMap<String, String>,
) {
for region_def in &blueprint.context_layout.regions {
let mut region = Region::new(
region_def.name.clone(),
region_def.kind.clone(),
region_def.max_tokens,
);
region.summarizable = region_def.summarizable;
region.admission = region_def.admission;
window.add_region(region);
}
if window.get_region("tool_results").is_none() {
let tool_region = Region::new("tool_results".to_string(), RegionKind::Temporary, 5000);
window.add_region(tool_region);
}
if window.get_region("conversation").is_none() {
let conv_region = Region::new(
"conversation".to_string(),
RegionKind::SlidingWindow {
max_items: 50,
eviction_strategy: EvictionStrategy::PerItem,
},
10000,
);
window.add_region(conv_region);
}
if window
.get_region(crate::output_tool::FINAL_OUTPUT_REGION)
.is_none()
{
window.add_region(Region::new(
crate::output_tool::FINAL_OUTPUT_REGION.to_string(),
RegionKind::Pinned,
crate::output_tool::FINAL_OUTPUT_REGION_TOKENS,
));
}
for (name, content) in seeds {
let target = if name == "task" {
task_region_name(blueprint)
} else {
blueprint
.context_layout
.regions
.iter()
.find(|r| &r.name == name)
.map(|r| r.name.clone())
};
if let Some(region_name) = target {
let budget = window
.get_region(®ion_name)
.map(|r| r.max_tokens)
.unwrap_or(0);
let fitted = fit_seed_to_budget(content, budget);
let tokens = leviath_core::estimate_tokens(&fitted);
let _ = window.add_to_region(®ion_name, fitted, tokens);
}
}
}
const SEED_TRUNCATION_MARKER: &str =
"\n[...truncated by leviath: seed exceeded this region's budget]";
fn fit_seed_to_budget(content: &str, max_tokens: usize) -> String {
let allowed = max_tokens.saturating_sub(1).saturating_mul(4);
if content.len() <= allowed {
return content.to_string();
}
let Some(room) = allowed.checked_sub(SEED_TRUNCATION_MARKER.len()) else {
return String::new();
};
format!(
"{}{SEED_TRUNCATION_MARKER}",
truncate_at_boundary(content, room)
)
}
fn task_region_name(blueprint: &Blueprint) -> Option<String> {
blueprint
.context_layout
.regions
.iter()
.find(|r| r.name == "task" && matches!(r.kind, RegionKind::Pinned))
.or_else(|| {
blueprint
.context_layout
.regions
.iter()
.find(|r| matches!(r.kind, RegionKind::Pinned))
})
.map(|r| r.name.clone())
}
pub fn init_window(window: &mut ContextWindow, blueprint: &Blueprint, task: &str) {
let seeds = HashMap::from([("task".to_string(), task.to_string())]);
init_window_seeded(window, blueprint, &seeds);
}
pub fn apply_layout(window: &mut ContextWindow, layout: &ContextLayout) {
let mut new_regions = Vec::new();
let mut kept: std::collections::HashSet<&str> = std::collections::HashSet::new();
for region_def in &layout.regions {
let mut new_region = Region::new(
region_def.name.clone(),
region_def.kind.clone(),
region_def.max_tokens,
);
new_region.summarizable = region_def.summarizable;
new_region.admission = region_def.admission;
if let Some(existing) = window.get_region(®ion_def.name) {
for entry in &existing.content {
let _ = new_region.carry_entry(entry.clone());
}
new_region.taint = existing.taint.clone();
}
kept.insert(region_def.name.as_str());
new_regions.push(new_region);
}
let always_visible = [
"conversation",
"tool_results",
crate::output_tool::FINAL_OUTPUT_REGION,
leviath_core::layout::STAGE_INSTRUCTIONS_REGION,
];
let mut hidden = std::collections::HashSet::new();
for existing in &window.regions {
if kept.contains(existing.name.as_str()) {
continue;
}
let mut carried = Region::new(
existing.name.clone(),
existing.kind.clone(),
existing.max_tokens,
);
carried.summarizable = existing.summarizable;
carried.admission = existing.admission;
for entry in &existing.content {
let _ = carried.carry_entry(entry.clone());
}
carried.taint = existing.taint.clone();
if !always_visible.contains(&existing.name.as_str()) {
hidden.insert(existing.name.clone());
}
new_regions.push(carried);
}
window.hidden = hidden;
window.regions = new_regions;
window.current_tokens = window.calculate_tokens();
}
pub fn ensure_stage_instructions_region(window: &mut ContextWindow, prompts: &[Option<String>]) {
let declared = leviath_core::layout::STAGE_INSTRUCTIONS_REGION;
if window.get_region(declared).is_some() {
return;
}
let widest = prompts
.iter()
.flatten()
.map(|p| leviath_core::estimate_tokens(&format!("[Stage instructions: {p}]")))
.max();
let Some(widest) = widest.filter(|t| *t > 0) else {
return;
};
let ceiling = window.max_tokens / INSTRUCTIONS_SHARE_OF_WINDOW;
window.add_region(Region::new(
declared.to_string(),
RegionKind::Pinned,
widest.min(ceiling),
));
}
const INSTRUCTIONS_SHARE_OF_WINDOW: usize = 4;
#[cfg(test)]
mod tests {
use super::{
SEED_TRUNCATION_MARKER, apply_layout, fit_seed_to_budget, init_window, init_window_seeded,
};
use crate::ContextWindow;
use leviath_core::{
Blueprint, ContextLayout, EvictionStrategy, RegionKind, Stage, blueprint::ModelConfig,
layout::RegionDefinition,
};
use std::collections::HashMap;
fn blueprint_with(regions: Vec<RegionDefinition>) -> Blueprint {
let layout = ContextLayout::new(regions, 100_000);
let stages = vec![Stage::new(
"main".to_string(),
ModelConfig::new("anthropic".to_string(), "claude-sonnet-4".to_string()),
)];
Blueprint::new("bp".to_string(), "desc".to_string(), stages, layout)
}
fn seeded_window(bp: &Blueprint, task: &str) -> ContextWindow {
let mut window = ContextWindow::new(100_000);
init_window(&mut window, bp, task);
window
}
#[test]
fn a_layout_that_declares_final_output_keeps_its_own() {
const DECLARED_TOKENS: usize = 12_345;
let bp = blueprint_with(vec![
RegionDefinition::new("task".to_string(), RegionKind::Pinned, 1_000),
RegionDefinition::new(
crate::output_tool::FINAL_OUTPUT_REGION.to_string(),
RegionKind::Pinned,
DECLARED_TOKENS,
),
]);
let window = seeded_window(&bp, "t");
assert_eq!(
window
.get_region(crate::output_tool::FINAL_OUTPUT_REGION)
.expect("the region is there")
.max_tokens,
DECLARED_TOKENS,
"the blueprint's own budget survives"
);
}
#[test]
fn a_layout_without_final_output_gets_the_default_one() {
let bp = blueprint_with(vec![RegionDefinition::new(
"task".to_string(),
RegionKind::Pinned,
1_000,
)]);
let window = seeded_window(&bp, "t");
assert_eq!(
window
.get_region(crate::output_tool::FINAL_OUTPUT_REGION)
.expect("added for us")
.max_tokens,
crate::output_tool::FINAL_OUTPUT_REGION_TOKENS
);
}
#[test]
fn init_window_seeded_fills_multiple_named_regions_and_ignores_unknown() {
let bp = blueprint_with(vec![
RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
RegionDefinition::new("criteria".to_string(), RegionKind::Pinned, 5000),
]);
let seeds = HashMap::from([
("task".to_string(), "build a parser".to_string()),
("criteria".to_string(), "focus on safety".to_string()),
("ghost".to_string(), "no such region".to_string()),
]);
let mut window = ContextWindow::new(100_000);
init_window_seeded(&mut window, &bp, &seeds);
assert!(
window
.get_region("task")
.unwrap()
.content
.iter()
.any(|e| e.content.contains("build a parser"))
);
assert!(
window
.get_region("criteria")
.unwrap()
.content
.iter()
.any(|e| e.content.contains("focus on safety"))
);
assert!(window.get_region("ghost").is_none());
}
#[test]
fn fit_seed_to_budget_leaves_a_fitting_seed_untouched() {
assert_eq!(fit_seed_to_budget("hello", 100), "hello");
let exact = "x".repeat(36);
assert_eq!(fit_seed_to_budget(&exact, 10), exact);
}
fn estimated_tokens(fitted: &str) -> usize {
leviath_core::estimate_tokens(fitted)
}
#[test]
fn fit_seed_to_budget_truncates_and_marks_an_oversized_seed() {
let big = "x".repeat(10_000);
let fitted = fit_seed_to_budget(&big, 100);
assert!(fitted.ends_with(SEED_TRUNCATION_MARKER));
let estimate = estimated_tokens(&fitted);
assert!(estimate <= 100, "estimate was {estimate}");
}
#[test]
fn fit_seed_to_budget_cuts_on_a_char_boundary() {
const MAX_TOKENS: usize = 60;
let room = (MAX_TOKENS - 1) * 4 - SEED_TRUNCATION_MARKER.len();
let mut s = "a".repeat(room - 1);
s.push('é'); s.push_str(&"b".repeat(500));
assert!(!s.is_char_boundary(room), "test must straddle the cut");
let fitted = fit_seed_to_budget(&s, MAX_TOKENS);
assert!(fitted.ends_with(SEED_TRUNCATION_MARKER));
assert!(estimated_tokens(&fitted) <= MAX_TOKENS);
assert_eq!(
fitted,
format!("{}{SEED_TRUNCATION_MARKER}", "a".repeat(room - 1))
);
}
#[test]
fn fit_seed_to_budget_yields_nothing_when_even_the_marker_cannot_fit() {
assert_eq!(fit_seed_to_budget("some content here", 2), "");
assert_eq!(fit_seed_to_budget("x", 0), "");
}
#[test]
fn init_window_seeded_truncates_a_seed_larger_than_its_region() {
let bp = blueprint_with(vec![RegionDefinition::new(
"facts".to_string(),
RegionKind::Pinned,
50,
)]);
let seeds = HashMap::from([("facts".to_string(), "y".repeat(10_000))]);
let mut window = ContextWindow::new(100_000);
init_window_seeded(&mut window, &bp, &seeds);
let region = window.get_region("facts").unwrap();
assert!(
!region.content.is_empty(),
"an oversized seed must be trimmed, not dropped"
);
assert!(region.content[0].content.ends_with(SEED_TRUNCATION_MARKER));
}
#[test]
fn init_window_seeded_task_key_falls_back_to_first_pinned() {
let bp = blueprint_with(vec![RegionDefinition::new(
"system".to_string(),
RegionKind::Pinned,
5000,
)]);
let seeds = HashMap::from([("task".to_string(), "fallback text".to_string())]);
let mut window = ContextWindow::new(100_000);
init_window_seeded(&mut window, &bp, &seeds);
assert!(
window
.get_region("system")
.unwrap()
.content
.iter()
.any(|e| e.content.contains("fallback text"))
);
}
#[test]
fn init_prefers_named_task_region_and_keeps_existing_infra_regions() {
let bp = blueprint_with(vec![
RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
RegionDefinition::new("tool_results".to_string(), RegionKind::Temporary, 5000),
RegionDefinition::new(
"conversation".to_string(),
RegionKind::SlidingWindow {
max_items: 10,
eviction_strategy: EvictionStrategy::PerItem,
},
10_000,
),
]);
let window = seeded_window(&bp, "do the thing");
assert!(
window
.get_region("task")
.unwrap()
.content
.iter()
.any(|e| e.content.contains("do the thing"))
);
assert_eq!(
window
.regions
.iter()
.filter(|r| r.name == "tool_results")
.count(),
1
);
assert_eq!(
window
.regions
.iter()
.filter(|r| r.name == "conversation")
.count(),
1
);
}
#[test]
fn init_adds_infra_regions_and_falls_back_to_first_pinned() {
let bp = blueprint_with(vec![RegionDefinition::new(
"system".to_string(),
RegionKind::Pinned,
5000,
)]);
let window = seeded_window(&bp, "seed task");
assert!(window.get_region("tool_results").is_some());
assert!(window.get_region("conversation").is_some());
assert!(
window
.get_region("system")
.unwrap()
.content
.iter()
.any(|e| e.content.contains("seed task"))
);
}
#[test]
fn init_without_pinned_region_does_not_seed_task() {
let bp = blueprint_with(vec![RegionDefinition::new(
"scratch".to_string(),
RegionKind::Temporary,
5000,
)]);
let window = seeded_window(&bp, "unseeded task");
assert!(window.get_region("scratch").unwrap().content.is_empty());
assert!(window.get_region("tool_results").is_some());
assert!(window.get_region("conversation").is_some());
}
#[test]
fn init_task_named_region_that_is_not_pinned_falls_back_to_first_pinned() {
let bp = blueprint_with(vec![
RegionDefinition::new("task".to_string(), RegionKind::Temporary, 5000),
RegionDefinition::new("system".to_string(), RegionKind::Pinned, 5000),
]);
let window = seeded_window(&bp, "fallback seed");
assert!(window.get_region("task").unwrap().content.is_empty());
assert!(
window
.get_region("system")
.unwrap()
.content
.iter()
.any(|e| e.content.contains("fallback seed"))
);
}
#[test]
fn apply_layout_preserves_overlapping_content_and_creates_new_regions() {
let bp = blueprint_with(vec![RegionDefinition::new(
"system".to_string(),
RegionKind::Pinned,
5000,
)]);
let mut window = seeded_window(&bp, "carried content");
let new_layout = ContextLayout::new(
vec![
RegionDefinition::new("system".to_string(), RegionKind::Pinned, 5000),
RegionDefinition::new("scratch".to_string(), RegionKind::Temporary, 3000),
],
8000,
);
apply_layout(&mut window, &new_layout);
assert_eq!(window.regions.len(), 5);
assert!(window.get_region("conversation").is_some());
assert!(window.get_region("tool_results").is_some());
assert!(
window
.get_region(crate::output_tool::FINAL_OUTPUT_REGION)
.is_some(),
"a submitted answer must survive a stage transition"
);
assert!(
window
.get_region("system")
.unwrap()
.content
.iter()
.any(|e| e.content.contains("carried content"))
);
assert!(window.get_region("scratch").unwrap().content.is_empty());
assert_eq!(window.current_tokens, window.calculate_tokens());
assert!(window.current_tokens > 0);
}
#[test]
fn apply_layout_preserves_entry_kinds_and_taint_across_swap() {
let bp = blueprint_with(vec![RegionDefinition::new(
"task".to_string(),
RegionKind::Pinned,
5000,
)]);
let mut window = seeded_window(&bp, "the task");
window
.add_typed_entry(
"conversation",
leviath_core::EntryKind::AssistantTurn {
tool_calls: vec![leviath_core::SerializedToolCall {
id: "call_9".to_string(),
name: "shell".to_string(),
arguments: serde_json::json!({"command": "ls"}),
thought_signature: None,
}],
},
"running ls".to_string(),
10,
)
.unwrap();
window
.add_typed_entry(
"conversation",
leviath_core::EntryKind::ToolResult {
tool_call_id: "call_9".to_string(),
tool_name: "shell".to_string(),
is_error: false,
},
"file_a\nfile_b".to_string(),
10,
)
.unwrap();
window
.get_region_mut("conversation")
.unwrap()
.enable_taint_tracking();
let omitting = ContextLayout::new(
vec![RegionDefinition::new(
"task".to_string(),
RegionKind::Pinned,
5000,
)],
8000,
);
apply_layout(&mut window, &omitting);
let declaring = ContextLayout::new(
vec![
RegionDefinition::new("task".to_string(), RegionKind::Pinned, 5000),
RegionDefinition::new(
"conversation".to_string(),
RegionKind::SlidingWindow {
max_items: 10,
eviction_strategy: EvictionStrategy::PerItem,
},
10_000,
),
],
20_000,
);
apply_layout(&mut window, &declaring);
let conv = window.get_region("conversation").unwrap();
assert!(
conv.content.iter().any(|e| matches!(
&e.kind,
leviath_core::EntryKind::AssistantTurn { tool_calls }
if tool_calls.iter().any(|c| c.id == "call_9")
)),
"assistant turn must keep its typed tool_calls through both carry paths"
);
assert!(
conv.content.iter().any(|e| matches!(
&e.kind,
leviath_core::EntryKind::ToolResult { tool_call_id, .. }
if tool_call_id == "call_9"
)),
"tool result must keep its typed pairing through both carry paths"
);
assert!(
conv.taint.is_some(),
"region-level taint state must carry across layout swaps"
);
}
#[test]
fn apply_layout_carries_conversation_when_new_layout_omits_it() {
let bp = blueprint_with(vec![RegionDefinition::new(
"task".to_string(),
RegionKind::Pinned,
5000,
)]);
let mut window = seeded_window(&bp, "the task");
window
.add_typed_entry(
"conversation",
leviath_core::EntryKind::UserMessage,
"hello from stage 0".to_string(),
10,
)
.unwrap();
let next = ContextLayout::new(
vec![RegionDefinition::new(
"task".to_string(),
RegionKind::Pinned,
5000,
)],
8000,
);
apply_layout(&mut window, &next);
let conv = window
.get_region("conversation")
.expect("conversation carried across transition");
assert!(
conv.content
.iter()
.any(|e| e.content.contains("hello from stage 0")),
"carried conversation must retain its history"
);
}
}