enact_context/

window.rs

//! Context Window Management
//!
//! The ContextWindow manages all segments within token budget limits.
//!
//! @see packages/enact-schemas/src/context.schemas.ts

use crate::budget::{BudgetHealth, ContextBudget};
use crate::compactor::{CompactionResult, CompactionStrategyType, Compactor};
use crate::segment::{ContextSegment, ContextSegmentType};
use crate::token_counter::TokenCounter;
use chrono::{DateTime, Utc};
use enact_core::kernel::ExecutionId;
use serde::{Deserialize, Serialize};
use std::time::Instant;
use thiserror::Error;

/// Context window errors
#[derive(Debug, Error)]
pub enum ContextWindowError {
    #[error("Token counter error: {0}")]
    TokenCounter(String),

    #[error("Budget exceeded: need {needed} tokens, only {available} available")]
    BudgetExceeded { needed: usize, available: usize },

    #[error("Segment budget exceeded for {segment_type:?}: need {needed}, max {max}")]
    SegmentBudgetExceeded {
        segment_type: ContextSegmentType,
        needed: usize,
        max: usize,
    },

    #[error("Compaction failed: {0}")]
    CompactionFailed(String),
}

/// Context window state
///
/// Matches `contextWindowStateSchema` in @enact/schemas
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ContextWindowState {
    /// Execution ID
    pub execution_id: ExecutionId,

    /// All segments currently in context
    pub segments: Vec<ContextSegment>,

    /// Current budget state
    pub budget: ContextBudget,

    /// Compaction history
    pub compaction_history: Vec<CompactionResult>,

    /// Number of compactions performed
    pub compaction_count: u32,

    /// Total tokens saved by compaction
    pub total_tokens_saved: usize,

    /// Current health status
    pub health: BudgetHealth,

    /// Last updated timestamp
    pub updated_at: DateTime<Utc>,
}

/// Context Window - manages segments within budget
pub struct ContextWindow {
    /// Execution ID
    execution_id: ExecutionId,

    /// All segments in the context
    segments: Vec<ContextSegment>,

    /// Token budget
    budget: ContextBudget,

    /// Token counter
    token_counter: TokenCounter,

    /// Compaction history
    compaction_history: Vec<CompactionResult>,

    /// Next sequence number
    next_sequence: u64,
}

impl ContextWindow {
    /// Create a new context window
    pub fn new(budget: ContextBudget) -> Result<Self, ContextWindowError> {
        let token_counter =
            TokenCounter::new().map_err(|e| ContextWindowError::TokenCounter(e.to_string()))?;

        Ok(Self {
            execution_id: budget.execution_id.clone(),
            segments: Vec::new(),
            budget,
            token_counter,
            compaction_history: Vec::new(),
            next_sequence: 0,
        })
    }

    /// Create with a preset budget
    pub fn with_preset_gpt4_128k(execution_id: ExecutionId) -> Result<Self, ContextWindowError> {
        Self::new(ContextBudget::preset_gpt4_128k(execution_id))
    }

    /// Create with Claude 200K preset
    pub fn with_preset_claude_200k(execution_id: ExecutionId) -> Result<Self, ContextWindowError> {
        Self::new(ContextBudget::preset_claude_200k(execution_id))
    }

    /// Create with default (8K) preset
    pub fn with_preset_default(execution_id: ExecutionId) -> Result<Self, ContextWindowError> {
        Self::new(ContextBudget::preset_default(execution_id))
    }

    /// Get the execution ID
    pub fn execution_id(&self) -> &ExecutionId {
        &self.execution_id
    }

    /// Get all segments
    pub fn segments(&self) -> &[ContextSegment] {
        &self.segments
    }

    /// Get the budget
    pub fn budget(&self) -> &ContextBudget {
        &self.budget
    }

    /// Get mutable budget
    pub fn budget_mut(&mut self) -> &mut ContextBudget {
        &mut self.budget
    }

    /// Count tokens for text
    pub fn count_tokens(&self, text: &str) -> usize {
        self.token_counter.count(text)
    }

    /// Add a segment with automatic token counting
    pub fn add_segment_auto(
        &mut self,
        mut segment: ContextSegment,
    ) -> Result<(), ContextWindowError> {
        // Count tokens if not already set
        if segment.token_count == 0 {
            segment.token_count = self.token_counter.count(&segment.content);
        }

        self.add_segment(segment)
    }

    /// Add a segment to the context window
    pub fn add_segment(&mut self, mut segment: ContextSegment) -> Result<(), ContextWindowError> {
        // Check segment budget
        if let Some(seg_budget) = self.budget.get_segment(segment.segment_type) {
            let new_usage = seg_budget.current_tokens + segment.token_count;
            if new_usage > seg_budget.max_tokens {
                return Err(ContextWindowError::SegmentBudgetExceeded {
                    segment_type: segment.segment_type,
                    needed: segment.token_count,
                    max: seg_budget.max_tokens - seg_budget.current_tokens,
                });
            }
        }

        // Check total budget
        let new_total = self.budget.used_tokens + segment.token_count;
        if new_total > self.budget.available_tokens {
            return Err(ContextWindowError::BudgetExceeded {
                needed: segment.token_count,
                available: self.budget.remaining(),
            });
        }

        // Assign sequence number
        segment.sequence = self.next_sequence;
        self.next_sequence += 1;

        // Update budget
        self.budget
            .add_tokens(segment.segment_type, segment.token_count);

        // Add segment
        self.segments.push(segment);

        Ok(())
    }

    /// Remove a segment by ID
    pub fn remove_segment(&mut self, segment_id: &str) -> bool {
        if let Some(pos) = self.segments.iter().position(|s| s.id == segment_id) {
            let segment = self.segments.remove(pos);
            self.budget
                .remove_tokens(segment.segment_type, segment.token_count);
            true
        } else {
            false
        }
    }

    /// Get segments of a specific type
    pub fn segments_of_type(&self, segment_type: ContextSegmentType) -> Vec<&ContextSegment> {
        self.segments
            .iter()
            .filter(|s| s.segment_type == segment_type)
            .collect()
    }

    /// Total tokens currently used
    pub fn used_tokens(&self) -> usize {
        self.budget.used_tokens
    }

    /// Remaining tokens available
    pub fn remaining_tokens(&self) -> usize {
        self.budget.remaining()
    }

    /// Check if the window needs compaction
    pub fn needs_compaction(&self) -> bool {
        self.budget.is_warning()
    }

    /// Check if the window is in critical state
    pub fn is_critical(&self) -> bool {
        self.budget.is_critical()
    }

    /// Get health status
    pub fn health(&self) -> BudgetHealth {
        self.budget.health()
    }

    /// Compact the context using the given compactor
    pub fn compact(
        &mut self,
        compactor: &Compactor,
    ) -> Result<CompactionResult, ContextWindowError> {
        let start = Instant::now();
        let tokens_before = self.budget.used_tokens;

        let result = match compactor.strategy().strategy_type {
            CompactionStrategyType::Truncate => {
                compactor.compact_truncate(&mut self.segments, tokens_before)
            }
            CompactionStrategyType::SlidingWindow => {
                compactor.compact_sliding_window(&mut self.segments)
            }
            _ => {
                // Other strategies not implemented yet
                return Err(ContextWindowError::CompactionFailed(format!(
                    "Strategy {:?} not implemented",
                    compactor.strategy().strategy_type
                )));
            }
        };

        let duration_ms = start.elapsed().as_millis() as u64;

        match result {
            Ok(tokens_removed) => {
                // Recalculate budget
                self.recalculate_budget();

                let tokens_after = self.budget.used_tokens;
                let segments_compacted = (tokens_removed > 0) as usize;

                let compaction_result = CompactionResult::success(
                    self.execution_id.clone(),
                    compactor.strategy().strategy_type,
                    tokens_before,
                    tokens_after,
                    segments_compacted,
                    duration_ms,
                );

                self.compaction_history.push(compaction_result.clone());
                Ok(compaction_result)
            }
            Err(e) => {
                let compaction_result = CompactionResult::failure(
                    self.execution_id.clone(),
                    compactor.strategy().strategy_type,
                    tokens_before,
                    e.to_string(),
                    duration_ms,
                );

                self.compaction_history.push(compaction_result.clone());
                Err(ContextWindowError::CompactionFailed(e.to_string()))
            }
        }
    }

    /// Recalculate budget from current segments
    fn recalculate_budget(&mut self) {
        // Reset all segment budgets
        for seg_budget in &mut self.budget.segments {
            seg_budget.current_tokens = 0;
        }

        // Sum up tokens by segment type
        for segment in &self.segments {
            self.budget
                .add_tokens(segment.segment_type, segment.token_count);
        }
    }

    /// Build the context as a single string (for LLM calls)
    pub fn build_context(&self) -> String {
        let mut parts: Vec<&str> = Vec::new();

        // Sort segments by sequence
        let mut sorted: Vec<&ContextSegment> = self.segments.iter().collect();
        sorted.sort_by_key(|s| s.sequence);

        for segment in sorted {
            parts.push(&segment.content);
        }

        parts.join("\n\n")
    }

    /// Get the current state (for serialization)
    pub fn state(&self) -> ContextWindowState {
        ContextWindowState {
            execution_id: self.execution_id.clone(),
            segments: self.segments.clone(),
            budget: self.budget.clone(),
            compaction_history: self.compaction_history.clone(),
            compaction_count: self.compaction_history.len() as u32,
            total_tokens_saved: self.compaction_history.iter().map(|r| r.tokens_saved).sum(),
            health: self.budget.health(),
            updated_at: Utc::now(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_execution_id() -> ExecutionId {
        ExecutionId::new()
    }

    #[test]
    fn test_create_window() {
        let budget = ContextBudget::preset_default(test_execution_id());
        let window = ContextWindow::new(budget).unwrap();

        assert_eq!(window.used_tokens(), 0);
        assert!(window.remaining_tokens() > 0);
    }

    #[test]
    fn test_add_segment() {
        let budget = ContextBudget::preset_default(test_execution_id());
        let mut window = ContextWindow::new(budget).unwrap();

        let segment = ContextSegment::system("You are a helpful assistant.", 10);
        window.add_segment(segment).unwrap();

        assert_eq!(window.segments().len(), 1);
        assert_eq!(window.used_tokens(), 10);
    }

    #[test]
    fn test_health_tracking() {
        let budget = ContextBudget::preset_default(test_execution_id());
        let window = ContextWindow::new(budget).unwrap();

        assert_eq!(window.health(), BudgetHealth::Healthy);
        assert!(!window.needs_compaction());
    }
}