1use compact_str::CompactString;
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
7#[serde(rename_all = "snake_case")]
8pub enum ContextSectionPartition {
9 System,
10 Skill,
11 Memory,
12 Working,
13 History,
14 Artifacts,
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum SectionCachePolicy {
21 Static,
23 SessionCached,
25 TurnDynamic,
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
31#[serde(rename_all = "snake_case")]
32pub enum SectionInvalidation {
33 Never,
34 OnCompact,
35 OnSkillChange,
36 OnCapabilityChange,
37 OnMemoryRefresh,
38 EveryTurn,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct ContextSection {
44 pub id: CompactString,
45 pub partition: ContextSectionPartition,
46 pub priority: i16,
48 pub cache_policy: SectionCachePolicy,
49 pub invalidation: SectionInvalidation,
50 #[serde(default, skip_serializing_if = "Option::is_none")]
51 pub token_budget: Option<u32>,
52 #[serde(default)]
53 pub enabled: bool,
54 #[serde(default)]
56 pub is_pinned: bool,
57}
58
59impl ContextSection {
60 pub fn new(
61 id: impl Into<CompactString>,
62 partition: ContextSectionPartition,
63 priority: i16,
64 ) -> Self {
65 Self {
66 id: id.into(),
67 partition,
68 priority,
69 cache_policy: SectionCachePolicy::TurnDynamic,
70 invalidation: SectionInvalidation::EveryTurn,
71 token_budget: None,
72 enabled: true,
73 is_pinned: false,
74 }
75 }
76
77 pub fn pinned(mut self) -> Self {
78 self.is_pinned = true;
79 self
80 }
81
82 pub fn with_cache_policy(mut self, policy: SectionCachePolicy) -> Self {
83 self.cache_policy = policy;
84 self
85 }
86
87 pub fn with_invalidation(mut self, invalidation: SectionInvalidation) -> Self {
88 self.invalidation = invalidation;
89 self
90 }
91
92 pub fn with_token_budget(mut self, token_budget: u32) -> Self {
93 self.token_budget = Some(token_budget);
94 self
95 }
96
97 pub fn disabled(mut self) -> Self {
98 self.enabled = false;
99 self
100 }
101}
102
103#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct ContextSectionPlan {
106 pub ids: Vec<CompactString>,
107}
108
109#[derive(Debug, Clone, Default, Serialize, Deserialize)]
111pub struct ContextSectionRegistry {
112 sections: BTreeMap<CompactString, ContextSection>,
113}
114
115impl ContextSectionRegistry {
116 pub fn new() -> Self {
117 Self {
118 sections: BTreeMap::new(),
119 }
120 }
121
122 pub fn default_agent_sections() -> Self {
124 let mut registry = Self::new();
125 registry.upsert(
126 ContextSection::new("system.base", ContextSectionPartition::System, 100)
127 .with_cache_policy(SectionCachePolicy::Static)
128 .with_invalidation(SectionInvalidation::Never),
129 );
130 registry.upsert(
131 ContextSection::new("system.task_state", ContextSectionPartition::System, 90)
132 .with_cache_policy(SectionCachePolicy::TurnDynamic)
133 .with_invalidation(SectionInvalidation::EveryTurn),
134 );
135 registry.upsert(
136 ContextSection::new(
137 "capabilities.inventory",
138 ContextSectionPartition::System,
139 80,
140 )
141 .with_cache_policy(SectionCachePolicy::SessionCached)
142 .with_invalidation(SectionInvalidation::OnCapabilityChange),
143 );
144 registry.upsert(
145 ContextSection::new("skill.active", ContextSectionPartition::Skill, 70)
146 .with_cache_policy(SectionCachePolicy::SessionCached)
147 .with_invalidation(SectionInvalidation::OnSkillChange),
148 );
149 registry.upsert(
150 ContextSection::new("memory.retrieved", ContextSectionPartition::Memory, 60)
151 .with_cache_policy(SectionCachePolicy::TurnDynamic)
152 .with_invalidation(SectionInvalidation::OnMemoryRefresh),
153 );
154 registry.upsert(
155 ContextSection::new("working.signals", ContextSectionPartition::Working, 50)
156 .with_cache_policy(SectionCachePolicy::TurnDynamic)
157 .with_invalidation(SectionInvalidation::EveryTurn),
158 );
159 registry.upsert(
160 ContextSection::new("artifacts.references", ContextSectionPartition::Artifacts, 40)
161 .with_cache_policy(SectionCachePolicy::TurnDynamic)
162 .with_invalidation(SectionInvalidation::EveryTurn),
163 );
164 registry.upsert(
165 ContextSection::new("history.rolling", ContextSectionPartition::History, 10)
166 .with_cache_policy(SectionCachePolicy::TurnDynamic)
167 .with_invalidation(SectionInvalidation::OnCompact),
168 );
169 registry
170 }
171
172 pub fn upsert(&mut self, section: ContextSection) {
173 self.sections.insert(section.id.clone(), section);
174 }
175
176 pub fn get(&self, id: &str) -> Option<&ContextSection> {
177 self.sections.get(id)
178 }
179
180 pub fn len(&self) -> usize {
181 self.sections.len()
182 }
183
184 pub fn is_empty(&self) -> bool {
185 self.sections.is_empty()
186 }
187
188 pub fn sections(&self) -> Vec<&ContextSection> {
189 let mut sections = self.sections.values().collect::<Vec<_>>();
190 sections.sort_by(|a, b| b.priority.cmp(&a.priority).then_with(|| a.id.cmp(&b.id)));
191 sections
192 }
193
194 pub fn plan(&self) -> ContextSectionPlan {
195 ContextSectionPlan {
196 ids: self
197 .sections()
198 .into_iter()
199 .filter(|s| s.enabled)
200 .map(|s| s.id.clone())
201 .collect(),
202 }
203 }
204
205 pub fn pin(&mut self, id: &str) -> bool {
208 if let Some(section) = self.sections.get_mut(id) {
209 section.is_pinned = true;
210 true
211 } else {
212 false
213 }
214 }
215
216 pub fn unpin(&mut self, id: &str) -> bool {
219 if let Some(section) = self.sections.get_mut(id) {
220 section.is_pinned = false;
221 true
222 } else {
223 false
224 }
225 }
226
227 pub fn is_partition_pinned(&self, partition: ContextSectionPartition) -> bool {
229 self.sections
230 .values()
231 .any(|s| s.partition == partition && s.is_pinned)
232 }
233
234 pub fn invalidate(&mut self, event: SectionInvalidation) -> Vec<CompactString> {
236 let mut invalidated = Vec::new();
237 for section in self.sections.values_mut() {
238 let matches_event = section.invalidation == event
239 || section.invalidation == SectionInvalidation::EveryTurn
240 || event == SectionInvalidation::EveryTurn;
241 if matches_event && section.invalidation != SectionInvalidation::Never {
242 section.enabled = false;
243 invalidated.push(section.id.clone());
244 }
245 }
246 invalidated.sort();
247 invalidated
248 }
249}
250
251#[cfg(test)]
252mod tests {
253 use super::*;
254
255 #[test]
256 fn default_sections_include_capability_inventory() {
257 let registry = ContextSectionRegistry::default_agent_sections();
258
259 assert!(registry.get("capabilities.inventory").is_some());
260 assert!(registry.get("system.base").is_some());
261 }
262
263 #[test]
264 fn plan_orders_by_priority_then_id() {
265 let mut registry = ContextSectionRegistry::new();
266 registry.upsert(ContextSection::new(
267 "b",
268 ContextSectionPartition::System,
269 10,
270 ));
271 registry.upsert(ContextSection::new(
272 "a",
273 ContextSectionPartition::System,
274 10,
275 ));
276 registry.upsert(ContextSection::new(
277 "top",
278 ContextSectionPartition::System,
279 20,
280 ));
281
282 let ids = registry
283 .plan()
284 .ids
285 .into_iter()
286 .map(|id| id.to_string())
287 .collect::<Vec<_>>();
288
289 assert_eq!(ids, ["top", "a", "b"]);
290 }
291
292 #[test]
293 fn upsert_replaces_section() {
294 let mut registry = ContextSectionRegistry::new();
295 registry.upsert(ContextSection::new(
296 "same",
297 ContextSectionPartition::System,
298 1,
299 ));
300 registry.upsert(ContextSection::new(
301 "same",
302 ContextSectionPartition::Memory,
303 2,
304 ));
305
306 assert_eq!(registry.len(), 1);
307 assert_eq!(
308 registry.get("same").unwrap().partition,
309 ContextSectionPartition::Memory
310 );
311 }
312
313 #[test]
314 fn pin_marks_section_and_is_detected_by_partition() {
315 let mut registry = ContextSectionRegistry::default_agent_sections();
316
317 assert!(!registry.is_partition_pinned(ContextSectionPartition::History));
318 let found = registry.pin("history.rolling");
319 assert!(found);
320 assert!(registry.is_partition_pinned(ContextSectionPartition::History));
321 assert!(!registry.is_partition_pinned(ContextSectionPartition::System));
323 }
324
325 #[test]
326 fn unpin_restores_compressibility() {
327 let mut registry = ContextSectionRegistry::default_agent_sections();
328 registry.pin("history.rolling");
329 assert!(registry.is_partition_pinned(ContextSectionPartition::History));
330 let found = registry.unpin("history.rolling");
331 assert!(found);
332 assert!(!registry.is_partition_pinned(ContextSectionPartition::History));
333 }
334
335 #[test]
336 fn pin_returns_false_for_unknown_section() {
337 let mut registry = ContextSectionRegistry::new();
338 assert!(!registry.pin("nonexistent"));
339 assert!(!registry.unpin("nonexistent"));
340 }
341
342 #[test]
343 fn pinned_builder_sets_flag() {
344 let section = ContextSection::new("h", ContextSectionPartition::History, 10).pinned();
345 assert!(section.is_pinned);
346 }
347
348 #[test]
349 fn invalidation_disables_matching_sections() {
350 let mut registry = ContextSectionRegistry::new();
351 registry.upsert(
352 ContextSection::new("cap", ContextSectionPartition::System, 1)
353 .with_invalidation(SectionInvalidation::OnCapabilityChange),
354 );
355 registry.upsert(
356 ContextSection::new("static", ContextSectionPartition::System, 2)
357 .with_invalidation(SectionInvalidation::Never),
358 );
359
360 let invalidated = registry.invalidate(SectionInvalidation::OnCapabilityChange);
361
362 assert_eq!(invalidated, [CompactString::new("cap")]);
363 assert!(!registry.get("cap").unwrap().enabled);
364 assert!(registry.get("static").unwrap().enabled);
365 }
366}