enact-context 0.0.2

//! Context Segment Types
//!
//! Segments are portions of the context window with their own priority and compressibility.
//!
//! @see packages/enact-schemas/src/context.schemas.ts

use chrono::{DateTime, Utc};
use enact_core::kernel::StepId;
use serde::{Deserialize, Serialize};

/// Types of segments in the context window
///
/// Matches `contextSegmentTypeSchema` in @enact/schemas
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContextSegmentType {
    /// System prompt (typically not compressible)
    System,
    /// Conversation history (compressible)
    History,
    /// Current working context (partially compressible)
    WorkingMemory,
    /// Tool execution results (compressible)
    ToolResults,
    /// Retrieved RAG content (compressible)
    RagContext,
    /// Current user input (not compressible)
    UserInput,
    /// Agent's thinking/scratchpad (compressible)
    AgentScratchpad,
    /// Summary from child execution (not compressible)
    ChildSummary,
    /// Mid-execution guidance from inbox (not compressible)
    Guidance,
}

impl ContextSegmentType {
    /// Whether this segment type is compressible by default
    pub fn is_compressible(&self) -> bool {
        match self {
            Self::System => false,
            Self::History => true,
            Self::WorkingMemory => true,
            Self::ToolResults => true,
            Self::RagContext => true,
            Self::UserInput => false,
            Self::AgentScratchpad => true,
            Self::ChildSummary => false,
            Self::Guidance => false,
        }
    }

    /// Default priority for this segment type
    pub fn default_priority(&self) -> ContextPriority {
        match self {
            Self::System => ContextPriority::Critical,
            Self::UserInput => ContextPriority::Critical,
            Self::Guidance => ContextPriority::High,
            Self::ChildSummary => ContextPriority::High,
            Self::History => ContextPriority::Medium,
            Self::WorkingMemory => ContextPriority::Medium,
            Self::ToolResults => ContextPriority::Medium,
            Self::RagContext => ContextPriority::Low,
            Self::AgentScratchpad => ContextPriority::Low,
        }
    }
}

/// Priority levels for context segments
///
/// Higher priority segments are preserved longer during compaction.
/// Matches `contextPrioritySchema` in @enact/schemas
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContextPriority {
    /// Compress first (old RAG results, agent scratchpad)
    Low = 0,
    /// Normal compression (older history, tool results)
    Medium = 1,
    /// Compress last (recent history, guidance)
    High = 2,
    /// Never compress (system prompt, current input)
    Critical = 3,
}

/// A portion of the context window
///
/// Matches `contextSegmentSchema` in @enact/schemas
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ContextSegment {
    /// Unique segment identifier
    pub id: String,

    /// Segment type
    #[serde(rename = "type")]
    pub segment_type: ContextSegmentType,

    /// Segment content
    pub content: String,

    /// Token count for this segment
    pub token_count: usize,

    /// Priority for compaction decisions
    pub priority: ContextPriority,

    /// Whether this segment can be compressed
    pub compressible: bool,

    /// Source step (for tool results, child summaries)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source_step_id: Option<StepId>,

    /// Timestamp when segment was added
    pub added_at: DateTime<Utc>,

    /// Sequence number for ordering
    pub sequence: u64,
}

impl ContextSegment {
    /// Create a new context segment
    pub fn new(
        segment_type: ContextSegmentType,
        content: String,
        token_count: usize,
        sequence: u64,
    ) -> Self {
        Self {
            id: format!("seg_{}", uuid::Uuid::new_v4()),
            segment_type,
            content,
            token_count,
            priority: segment_type.default_priority(),
            compressible: segment_type.is_compressible(),
            source_step_id: None,
            added_at: Utc::now(),
            sequence,
        }
    }

    /// Create a system prompt segment
    pub fn system(content: impl Into<String>, token_count: usize) -> Self {
        Self::new(ContextSegmentType::System, content.into(), token_count, 0)
    }

    /// Create a user input segment
    pub fn user_input(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
        Self::new(
            ContextSegmentType::UserInput,
            content.into(),
            token_count,
            sequence,
        )
    }

    /// Create a history segment
    pub fn history(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
        Self::new(
            ContextSegmentType::History,
            content.into(),
            token_count,
            sequence,
        )
    }

    /// Create a tool results segment
    pub fn tool_results(
        content: impl Into<String>,
        token_count: usize,
        sequence: u64,
        step_id: StepId,
    ) -> Self {
        let mut segment = Self::new(
            ContextSegmentType::ToolResults,
            content.into(),
            token_count,
            sequence,
        );
        segment.source_step_id = Some(step_id);
        segment
    }

    /// Create a RAG context segment
    pub fn rag_context(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
        Self::new(
            ContextSegmentType::RagContext,
            content.into(),
            token_count,
            sequence,
        )
    }

    /// Create a child summary segment
    pub fn child_summary(
        content: impl Into<String>,
        token_count: usize,
        sequence: u64,
        step_id: StepId,
    ) -> Self {
        let mut segment = Self::new(
            ContextSegmentType::ChildSummary,
            content.into(),
            token_count,
            sequence,
        );
        segment.source_step_id = Some(step_id);
        segment
    }

    /// Create a guidance segment (from inbox)
    pub fn guidance(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
        Self::new(
            ContextSegmentType::Guidance,
            content.into(),
            token_count,
            sequence,
        )
    }

    /// Set custom priority
    pub fn with_priority(mut self, priority: ContextPriority) -> Self {
        self.priority = priority;
        self
    }

    /// Mark as non-compressible
    pub fn non_compressible(mut self) -> Self {
        self.compressible = false;
        self
    }
}

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

    #[test]
    fn test_segment_type_compressibility() {
        assert!(!ContextSegmentType::System.is_compressible());
        assert!(ContextSegmentType::History.is_compressible());
        assert!(!ContextSegmentType::UserInput.is_compressible());
        assert!(ContextSegmentType::ToolResults.is_compressible());
    }

    #[test]
    fn test_segment_type_priority() {
        assert_eq!(
            ContextSegmentType::System.default_priority(),
            ContextPriority::Critical
        );
        assert_eq!(
            ContextSegmentType::History.default_priority(),
            ContextPriority::Medium
        );
        assert_eq!(
            ContextSegmentType::RagContext.default_priority(),
            ContextPriority::Low
        );
    }

    #[test]
    fn test_priority_ordering() {
        assert!(ContextPriority::Critical > ContextPriority::High);
        assert!(ContextPriority::High > ContextPriority::Medium);
        assert!(ContextPriority::Medium > ContextPriority::Low);
    }

    #[test]
    fn test_create_segment() {
        let segment = ContextSegment::system("You are helpful", 10);
        assert_eq!(segment.segment_type, ContextSegmentType::System);
        assert_eq!(segment.priority, ContextPriority::Critical);
        assert!(!segment.compressible);
    }
}