enya-agent 0.1.0

package/crates/enact-context/src/budget.rs

@@ -0,0 +1,314 @@
+//! Context Budget Management
+//!
+//! Token allocation across different segment types.
+//!
+//! @see packages/enact-schemas/src/context.schemas.ts
+
+use crate::segment::ContextSegmentType;
+use chrono::{DateTime, Utc};
+use enact_core::kernel::ExecutionId;
+use serde::{Deserialize, Serialize};
+
+/// Token allocation for a segment type
+///
+/// Matches `segmentBudgetSchema` in @enact/schemas
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SegmentBudget {
+    /// Segment type
+    #[serde(rename = "type")]
+    pub segment_type: ContextSegmentType,
+
+    /// Maximum tokens for this segment
+    pub max_tokens: usize,
+
+    /// Current token usage
+    pub current_tokens: usize,
+
+    /// Reserved tokens (cannot be used by other segments)
+    pub reserved_tokens: usize,
+
+    /// Whether this segment can borrow from others
+    pub can_borrow: bool,
+
+    /// Whether this segment can lend to others
+    pub can_lend: bool,
+}
+
+impl SegmentBudget {
+    /// Create a new segment budget
+    pub fn new(segment_type: ContextSegmentType, max_tokens: usize) -> Self {
+        Self {
+            segment_type,
+            max_tokens,
+            current_tokens: 0,
+            reserved_tokens: 0,
+            can_borrow: false,
+            can_lend: false,
+        }
+    }
+
+    /// Available tokens in this segment
+    pub fn available(&self) -> usize {
+        self.max_tokens.saturating_sub(self.current_tokens)
+    }
+
+    /// Usage percentage (0-100)
+    pub fn usage_percent(&self) -> u8 {
+        if self.max_tokens == 0 {
+            return 0;
+        }
+        ((self.current_tokens as f64 / self.max_tokens as f64) * 100.0) as u8
+    }
+
+    /// Whether this segment is over budget
+    pub fn is_over_budget(&self) -> bool {
+        self.current_tokens > self.max_tokens
+    }
+}
+
+/// Token allocation across all segments
+///
+/// Matches `contextBudgetSchema` in @enact/schemas
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ContextBudget {
+    /// Execution this budget belongs to
+    pub execution_id: ExecutionId,
+
+    /// Total context window size (model limit)
+    pub total_tokens: usize,
+
+    /// Tokens reserved for output generation
+    pub output_reserve: usize,
+
+    /// Available tokens for input context
+    pub available_tokens: usize,
+
+    /// Current total usage
+    pub used_tokens: usize,
+
+    /// Per-segment budgets
+    pub segments: Vec<SegmentBudget>,
+
+    /// Warning threshold (percentage, 0-100)
+    pub warning_threshold: u8,
+
+    /// Critical threshold (percentage, 0-100)
+    pub critical_threshold: u8,
+
+    /// Last updated timestamp
+    pub updated_at: DateTime<Utc>,
+}
+
+impl ContextBudget {
+    /// Create a new context budget
+    pub fn new(execution_id: ExecutionId, total_tokens: usize, output_reserve: usize) -> Self {
+        Self {
+            execution_id,
+            total_tokens,
+            output_reserve,
+            available_tokens: total_tokens.saturating_sub(output_reserve),
+            used_tokens: 0,
+            segments: Vec::new(),
+            warning_threshold: 80,
+            critical_threshold: 95,
+            updated_at: Utc::now(),
+        }
+    }
+
+    /// Create a GPT-4 128K preset
+    pub fn preset_gpt4_128k(execution_id: ExecutionId) -> Self {
+        let mut budget = Self::new(execution_id, 128_000, 4_096);
+        budget.segments = vec![
+            SegmentBudget::new(ContextSegmentType::System, 4_000),
+            SegmentBudget::new(ContextSegmentType::History, 60_000),
+            SegmentBudget::new(ContextSegmentType::WorkingMemory, 20_000),
+            SegmentBudget::new(ContextSegmentType::ToolResults, 20_000),
+            SegmentBudget::new(ContextSegmentType::RagContext, 15_000),
+            SegmentBudget::new(ContextSegmentType::UserInput, 2_000),
+            SegmentBudget::new(ContextSegmentType::AgentScratchpad, 2_000),
+            SegmentBudget::new(ContextSegmentType::ChildSummary, 500),
+            SegmentBudget::new(ContextSegmentType::Guidance, 500),
+        ];
+        budget
+    }
+
+    /// Create a GPT-4 32K preset
+    pub fn preset_gpt4_32k(execution_id: ExecutionId) -> Self {
+        let mut budget = Self::new(execution_id, 32_000, 2_048);
+        budget.segments = vec![
+            SegmentBudget::new(ContextSegmentType::System, 2_000),
+            SegmentBudget::new(ContextSegmentType::History, 15_000),
+            SegmentBudget::new(ContextSegmentType::WorkingMemory, 5_000),
+            SegmentBudget::new(ContextSegmentType::ToolResults, 4_000),
+            SegmentBudget::new(ContextSegmentType::RagContext, 3_000),
+            SegmentBudget::new(ContextSegmentType::UserInput, 1_000),
+            SegmentBudget::new(ContextSegmentType::AgentScratchpad, 500),
+            SegmentBudget::new(ContextSegmentType::ChildSummary, 250),
+            SegmentBudget::new(ContextSegmentType::Guidance, 250),
+        ];
+        budget
+    }
+
+    /// Create a Claude 200K preset
+    pub fn preset_claude_200k(execution_id: ExecutionId) -> Self {
+        let mut budget = Self::new(execution_id, 200_000, 4_096);
+        budget.segments = vec![
+            SegmentBudget::new(ContextSegmentType::System, 8_000),
+            SegmentBudget::new(ContextSegmentType::History, 100_000),
+            SegmentBudget::new(ContextSegmentType::WorkingMemory, 40_000),
+            SegmentBudget::new(ContextSegmentType::ToolResults, 25_000),
+            SegmentBudget::new(ContextSegmentType::RagContext, 15_000),
+            SegmentBudget::new(ContextSegmentType::UserInput, 4_000),
+            SegmentBudget::new(ContextSegmentType::AgentScratchpad, 2_000),
+            SegmentBudget::new(ContextSegmentType::ChildSummary, 1_000),
+            SegmentBudget::new(ContextSegmentType::Guidance, 1_000),
+        ];
+        budget
+    }
+
+    /// Create a default 8K preset
+    pub fn preset_default(execution_id: ExecutionId) -> Self {
+        let mut budget = Self::new(execution_id, 8_000, 1_024);
+        budget.segments = vec![
+            SegmentBudget::new(ContextSegmentType::System, 1_000),
+            SegmentBudget::new(ContextSegmentType::History, 3_000),
+            SegmentBudget::new(ContextSegmentType::WorkingMemory, 1_000),
+            SegmentBudget::new(ContextSegmentType::ToolResults, 1_000),
+            SegmentBudget::new(ContextSegmentType::RagContext, 500),
+            SegmentBudget::new(ContextSegmentType::UserInput, 500),
+            SegmentBudget::new(ContextSegmentType::AgentScratchpad, 0),
+            SegmentBudget::new(ContextSegmentType::ChildSummary, 0),
+            SegmentBudget::new(ContextSegmentType::Guidance, 0),
+        ];
+        budget
+    }
+
+    /// Get the budget for a specific segment type
+    pub fn get_segment(&self, segment_type: ContextSegmentType) -> Option<&SegmentBudget> {
+        self.segments.iter().find(|s| s.segment_type == segment_type)
+    }
+
+    /// Get mutable budget for a specific segment type
+    pub fn get_segment_mut(&mut self, segment_type: ContextSegmentType) -> Option<&mut SegmentBudget> {
+        self.segments.iter_mut().find(|s| s.segment_type == segment_type)
+    }
+
+    /// Update token usage for a segment
+    pub fn update_segment_usage(&mut self, segment_type: ContextSegmentType, tokens: usize) {
+        if let Some(segment) = self.get_segment_mut(segment_type) {
+            segment.current_tokens = tokens;
+        }
+        self.recalculate_total();
+    }
+
+    /// Add tokens to a segment
+    pub fn add_tokens(&mut self, segment_type: ContextSegmentType, tokens: usize) {
+        if let Some(segment) = self.get_segment_mut(segment_type) {
+            segment.current_tokens += tokens;
+        }
+        self.recalculate_total();
+    }
+
+    /// Remove tokens from a segment
+    pub fn remove_tokens(&mut self, segment_type: ContextSegmentType, tokens: usize) {
+        if let Some(segment) = self.get_segment_mut(segment_type) {
+            segment.current_tokens = segment.current_tokens.saturating_sub(tokens);
+        }
+        self.recalculate_total();
+    }
+
+    /// Recalculate total used tokens
+    fn recalculate_total(&mut self) {
+        self.used_tokens = self.segments.iter().map(|s| s.current_tokens).sum();
+        self.updated_at = Utc::now();
+    }
+
+    /// Remaining tokens
+    pub fn remaining(&self) -> usize {
+        self.available_tokens.saturating_sub(self.used_tokens)
+    }
+
+    /// Usage percentage (0-100)
+    pub fn usage_percent(&self) -> u8 {
+        if self.available_tokens == 0 {
+            return 0;
+        }
+        ((self.used_tokens as f64 / self.available_tokens as f64) * 100.0) as u8
+    }
+
+    /// Whether usage is at warning level
+    pub fn is_warning(&self) -> bool {
+        self.usage_percent() >= self.warning_threshold
+    }
+
+    /// Whether usage is at critical level
+    pub fn is_critical(&self) -> bool {
+        self.usage_percent() >= self.critical_threshold
+    }
+
+    /// Health status
+    pub fn health(&self) -> BudgetHealth {
+        if self.is_critical() {
+            BudgetHealth::Critical
+        } else if self.is_warning() {
+            BudgetHealth::Warning
+        } else {
+            BudgetHealth::Healthy
+        }
+    }
+}
+
+/// Budget health status
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum BudgetHealth {
+    Healthy,
+    Warning,
+    Critical,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_execution_id() -> ExecutionId {
+        ExecutionId::new()
+    }
+
+    #[test]
+    fn test_segment_budget() {
+        let mut budget = SegmentBudget::new(ContextSegmentType::History, 1000);
+        assert_eq!(budget.available(), 1000);
+        assert_eq!(budget.usage_percent(), 0);
+
+        budget.current_tokens = 500;
+        assert_eq!(budget.available(), 500);
+        assert_eq!(budget.usage_percent(), 50);
+    }
+
+    #[test]
+    fn test_context_budget_presets() {
+        let budget = ContextBudget::preset_gpt4_128k(test_execution_id());
+        assert_eq!(budget.total_tokens, 128_000);
+        assert_eq!(budget.output_reserve, 4_096);
+        assert!(!budget.segments.is_empty());
+    }
+
+    #[test]
+    fn test_budget_health() {
+        let mut budget = ContextBudget::preset_default(test_execution_id());
+
+        // Initially healthy
+        assert_eq!(budget.health(), BudgetHealth::Healthy);
+
+        // Add tokens to reach warning
+        budget.used_tokens = (budget.available_tokens as f64 * 0.85) as usize;
+        assert_eq!(budget.health(), BudgetHealth::Warning);
+
+        // Add more to reach critical
+        budget.used_tokens = (budget.available_tokens as f64 * 0.96) as usize;
+        assert_eq!(budget.health(), BudgetHealth::Critical);
+    }
+}