Skip to main content

deepstrike_core/context/
sections.rs

1use compact_str::CompactString;
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4
5/// Named context partition used by section planning.
6#[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/// Cache behavior for a section.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum SectionCachePolicy {
21    /// Stable across runs unless the application changes it.
22    Static,
23    /// Stable within a session until an invalidation event occurs.
24    SessionCached,
25    /// Recomputed on every turn.
26    TurnDynamic,
27}
28
29/// Events that can invalidate cached context sections.
30#[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/// One context section declaration.
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct ContextSection {
44    pub id: CompactString,
45    pub partition: ContextSectionPartition,
46    /// Higher priority sections are rendered earlier.
47    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    /// Pinned sections are exempt from GC/compression even under token pressure.
55    #[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/// Deterministic section plan produced by the registry.
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct ContextSectionPlan {
106    pub ids: Vec<CompactString>,
107}
108
109/// Registry for prompt/context sections and their lifecycle policy.
110#[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    /// Baseline sections that match DeepStrike's current 5-partition context.
123    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    /// Pin a section so its partition is exempt from GC compression.
206    /// Returns true if the section was found.
207    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    /// Unpin a section, allowing its partition to be compressed again.
217    /// Returns true if the section was found.
218    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    /// Returns true if any enabled section mapped to `partition` is pinned.
228    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    /// Mark sections invalidated by an event as disabled and return their ids.
235    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        // System partition unaffected
322        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}