Skip to main content

enact_context/
segment.rs

1//! Context Segment Types
2//!
3//! Segments are portions of the context window with their own priority and compressibility.
4//!
5//! @see packages/enact-schemas/src/context.schemas.ts
6
7use chrono::{DateTime, Utc};
8use enact_core::kernel::StepId;
9use serde::{Deserialize, Serialize};
10
11/// Types of segments in the context window
12///
13/// Matches `contextSegmentTypeSchema` in @enact/schemas
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
15#[serde(rename_all = "snake_case")]
16pub enum ContextSegmentType {
17    /// System prompt (typically not compressible)
18    System,
19    /// Conversation history (compressible)
20    History,
21    /// Current working context (partially compressible)
22    WorkingMemory,
23    /// Tool execution results (compressible)
24    ToolResults,
25    /// Retrieved RAG content (compressible)
26    RagContext,
27    /// Current user input (not compressible)
28    UserInput,
29    /// Agent's thinking/scratchpad (compressible)
30    AgentScratchpad,
31    /// Summary from child execution (not compressible)
32    ChildSummary,
33    /// Mid-execution guidance from inbox (not compressible)
34    Guidance,
35}
36
37impl ContextSegmentType {
38    /// Whether this segment type is compressible by default
39    pub fn is_compressible(&self) -> bool {
40        match self {
41            Self::System => false,
42            Self::History => true,
43            Self::WorkingMemory => true,
44            Self::ToolResults => true,
45            Self::RagContext => true,
46            Self::UserInput => false,
47            Self::AgentScratchpad => true,
48            Self::ChildSummary => false,
49            Self::Guidance => false,
50        }
51    }
52
53    /// Default priority for this segment type
54    pub fn default_priority(&self) -> ContextPriority {
55        match self {
56            Self::System => ContextPriority::Critical,
57            Self::UserInput => ContextPriority::Critical,
58            Self::Guidance => ContextPriority::High,
59            Self::ChildSummary => ContextPriority::High,
60            Self::History => ContextPriority::Medium,
61            Self::WorkingMemory => ContextPriority::Medium,
62            Self::ToolResults => ContextPriority::Medium,
63            Self::RagContext => ContextPriority::Low,
64            Self::AgentScratchpad => ContextPriority::Low,
65        }
66    }
67}
68
69/// Priority levels for context segments
70///
71/// Higher priority segments are preserved longer during compaction.
72/// Matches `contextPrioritySchema` in @enact/schemas
73#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
74#[serde(rename_all = "snake_case")]
75pub enum ContextPriority {
76    /// Compress first (old RAG results, agent scratchpad)
77    Low = 0,
78    /// Normal compression (older history, tool results)
79    Medium = 1,
80    /// Compress last (recent history, guidance)
81    High = 2,
82    /// Never compress (system prompt, current input)
83    Critical = 3,
84}
85
86/// A portion of the context window
87///
88/// Matches `contextSegmentSchema` in @enact/schemas
89#[derive(Debug, Clone, Serialize, Deserialize)]
90#[serde(rename_all = "camelCase")]
91pub struct ContextSegment {
92    /// Unique segment identifier
93    pub id: String,
94
95    /// Segment type
96    #[serde(rename = "type")]
97    pub segment_type: ContextSegmentType,
98
99    /// Segment content
100    pub content: String,
101
102    /// Token count for this segment
103    pub token_count: usize,
104
105    /// Priority for compaction decisions
106    pub priority: ContextPriority,
107
108    /// Whether this segment can be compressed
109    pub compressible: bool,
110
111    /// Source step (for tool results, child summaries)
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub source_step_id: Option<StepId>,
114
115    /// Timestamp when segment was added
116    pub added_at: DateTime<Utc>,
117
118    /// Sequence number for ordering
119    pub sequence: u64,
120}
121
122impl ContextSegment {
123    /// Create a new context segment
124    pub fn new(
125        segment_type: ContextSegmentType,
126        content: String,
127        token_count: usize,
128        sequence: u64,
129    ) -> Self {
130        Self {
131            id: format!("seg_{}", uuid::Uuid::new_v4()),
132            segment_type,
133            content,
134            token_count,
135            priority: segment_type.default_priority(),
136            compressible: segment_type.is_compressible(),
137            source_step_id: None,
138            added_at: Utc::now(),
139            sequence,
140        }
141    }
142
143    /// Create a system prompt segment
144    pub fn system(content: impl Into<String>, token_count: usize) -> Self {
145        Self::new(ContextSegmentType::System, content.into(), token_count, 0)
146    }
147
148    /// Create a user input segment
149    pub fn user_input(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
150        Self::new(
151            ContextSegmentType::UserInput,
152            content.into(),
153            token_count,
154            sequence,
155        )
156    }
157
158    /// Create a history segment
159    pub fn history(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
160        Self::new(
161            ContextSegmentType::History,
162            content.into(),
163            token_count,
164            sequence,
165        )
166    }
167
168    /// Create a tool results segment
169    pub fn tool_results(
170        content: impl Into<String>,
171        token_count: usize,
172        sequence: u64,
173        step_id: StepId,
174    ) -> Self {
175        let mut segment = Self::new(
176            ContextSegmentType::ToolResults,
177            content.into(),
178            token_count,
179            sequence,
180        );
181        segment.source_step_id = Some(step_id);
182        segment
183    }
184
185    /// Create a RAG context segment
186    pub fn rag_context(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
187        Self::new(
188            ContextSegmentType::RagContext,
189            content.into(),
190            token_count,
191            sequence,
192        )
193    }
194
195    /// Create a child summary segment
196    pub fn child_summary(
197        content: impl Into<String>,
198        token_count: usize,
199        sequence: u64,
200        step_id: StepId,
201    ) -> Self {
202        let mut segment = Self::new(
203            ContextSegmentType::ChildSummary,
204            content.into(),
205            token_count,
206            sequence,
207        );
208        segment.source_step_id = Some(step_id);
209        segment
210    }
211
212    /// Create a guidance segment (from inbox)
213    pub fn guidance(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
214        Self::new(
215            ContextSegmentType::Guidance,
216            content.into(),
217            token_count,
218            sequence,
219        )
220    }
221
222    /// Set custom priority
223    pub fn with_priority(mut self, priority: ContextPriority) -> Self {
224        self.priority = priority;
225        self
226    }
227
228    /// Mark as non-compressible
229    pub fn non_compressible(mut self) -> Self {
230        self.compressible = false;
231        self
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[test]
240    fn test_segment_type_compressibility() {
241        assert!(!ContextSegmentType::System.is_compressible());
242        assert!(ContextSegmentType::History.is_compressible());
243        assert!(!ContextSegmentType::UserInput.is_compressible());
244        assert!(ContextSegmentType::ToolResults.is_compressible());
245    }
246
247    #[test]
248    fn test_segment_type_priority() {
249        assert_eq!(
250            ContextSegmentType::System.default_priority(),
251            ContextPriority::Critical
252        );
253        assert_eq!(
254            ContextSegmentType::History.default_priority(),
255            ContextPriority::Medium
256        );
257        assert_eq!(
258            ContextSegmentType::RagContext.default_priority(),
259            ContextPriority::Low
260        );
261    }
262
263    #[test]
264    fn test_priority_ordering() {
265        assert!(ContextPriority::Critical > ContextPriority::High);
266        assert!(ContextPriority::High > ContextPriority::Medium);
267        assert!(ContextPriority::Medium > ContextPriority::Low);
268    }
269
270    #[test]
271    fn test_create_segment() {
272        let segment = ContextSegment::system("You are helpful", 10);
273        assert_eq!(segment.segment_type, ContextSegmentType::System);
274        assert_eq!(segment.priority, ContextPriority::Critical);
275        assert!(!segment.compressible);
276    }
277}