consciousness/src/agent/context.rs

// context.rs — Context window management
//
// Token counting, conversation trimming, and error classification.
// Journal entries are loaded from the memory graph store, not from
// a flat file — the parse functions are gone.

use std::sync::{Arc, RwLock};

use crate::agent::api::types::*;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use tiktoken_rs::CoreBPE;
use crate::agent::tools::working_stack;

/// A section of the context window, possibly with children.
#[derive(Debug, Clone)]
pub struct ContextSection {
    pub name: String,
    pub tokens: usize,
    pub content: String,
    pub children: Vec<ContextSection>,
}

/// Shared, live context state — agent writes, TUI reads for the debug screen.
pub type SharedContextState = Arc<RwLock<Vec<ContextSection>>>;

/// Create a new shared context state.
pub fn shared_context_state() -> SharedContextState {
    Arc::new(RwLock::new(Vec::new()))
}

/// A single journal entry with its timestamp and content.
#[derive(Debug, Clone)]
pub struct JournalEntry {
    pub timestamp: DateTime<Utc>,
    pub content: String,
}

/// Context window size in tokens (from config).
pub fn context_window() -> usize {
    crate::config::get().api_context_window
}

/// Context budget in tokens: 80% of the model's context window.
/// The remaining 20% is reserved for model output.
fn context_budget_tokens() -> usize {
    context_window() * 80 / 100
}

/// Dedup and trim conversation entries to fit within the context budget.
///
/// 1. Dedup: if the same memory key appears multiple times, keep only
///    the latest render (drop the earlier Memory entry and its
///    corresponding assistant tool_call message).
/// 2. Trim: drop oldest entries until the conversation fits, snapping
///    to user message boundaries.
pub fn trim_entries(
    context: &ContextState,
    entries: &[ConversationEntry],
    tokenizer: &CoreBPE,
) -> Vec<ConversationEntry> {
    let count = |s: &str| tokenizer.encode_with_special_tokens(s).len();

    // --- Phase 1: dedup memory entries by key (keep last) ---
    let mut seen_keys: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
    let mut drop_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();

    for (i, entry) in entries.iter().enumerate() {
        if let ConversationEntry::Memory { key, .. } = entry {
            if let Some(prev) = seen_keys.insert(key.as_str(), i) {
                drop_indices.insert(prev);
            }
        }
    }

    let deduped: Vec<ConversationEntry> = entries.iter().enumerate()
        .filter(|(i, _)| !drop_indices.contains(i))
        .map(|(_, e)| e.clone())
        .collect();

    // --- Phase 2: trim to fit context budget ---
    let max_tokens = context_budget_tokens();
    let identity_cost = count(&context.system_prompt)
        + context.personality.iter().map(|(_, c)| count(c)).sum::<usize>();
    let journal_cost: usize = context.journal.iter().map(|e| count(&e.content)).sum();
    let available = max_tokens
        .saturating_sub(identity_cost)
        .saturating_sub(journal_cost);

    let msg_costs: Vec<usize> = deduped.iter()
        .map(|e| msg_token_count(tokenizer, e.message())).collect();
    let total: usize = msg_costs.iter().sum();

    let mut skip = 0;
    let mut trimmed = total;
    while trimmed > available && skip < deduped.len() {
        trimmed -= msg_costs[skip];
        skip += 1;
    }

    // Walk forward to user message boundary
    while skip < deduped.len() && deduped[skip].message().role != Role::User {
        skip += 1;
    }

    deduped[skip..].to_vec()
}

/// Count the token footprint of a message using BPE tokenization.
pub fn msg_token_count(tokenizer: &CoreBPE, msg: &Message) -> usize {
    let count = |s: &str| tokenizer.encode_with_special_tokens(s).len();
    let content = msg.content.as_ref().map_or(0, |c| match c {
        MessageContent::Text(s) => count(s),
        MessageContent::Parts(parts) => parts.iter()
            .map(|p| match p {
                ContentPart::Text { text } => count(text),
                ContentPart::ImageUrl { .. } => 85,
            })
            .sum(),
    });
    let tools = msg.tool_calls.as_ref().map_or(0, |calls| {
        calls.iter()
            .map(|c| count(&c.function.arguments) + count(&c.function.name))
            .sum()
    });
    content + tools
}

/// Detect context window overflow errors from the API.
pub fn is_context_overflow(err: &anyhow::Error) -> bool {
    let msg = err.to_string().to_lowercase();
    msg.contains("context length")
        || msg.contains("token limit")
        || msg.contains("too many tokens")
        || msg.contains("maximum context")
        || msg.contains("prompt is too long")
        || msg.contains("request too large")
        || msg.contains("input validation error")
        || msg.contains("content length limit")
        || (msg.contains("400") && msg.contains("tokens"))
}

/// Detect model/provider errors delivered inside the SSE stream.
pub fn is_stream_error(err: &anyhow::Error) -> bool {
    err.to_string().contains("model stream error")
}

// --- Context state types ---

/// Conversation entry — either a regular message or memory content.
/// Memory entries preserve the original message for KV cache round-tripping.
#[derive(Debug, Clone)]
pub enum ConversationEntry {
    Message(Message),
    Memory { key: String, message: Message },
}

// Custom serde: serialize Memory with a "memory_key" field added to the message,
// plain messages serialize as-is. This keeps the conversation log readable.
impl Serialize for ConversationEntry {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        use serde::ser::SerializeMap;
        match self {
            Self::Message(m) => m.serialize(s),
            Self::Memory { key, message } => {
                let json = serde_json::to_value(message).map_err(serde::ser::Error::custom)?;
                let mut map = s.serialize_map(None)?;
                if let serde_json::Value::Object(obj) = json {
                    for (k, v) in obj {
                        map.serialize_entry(&k, &v)?;
                    }
                }
                map.serialize_entry("memory_key", key)?;
                map.end()
            }
        }
    }
}

impl<'de> Deserialize<'de> for ConversationEntry {
    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        let mut json: serde_json::Value = serde_json::Value::deserialize(d)?;
        if let Some(key) = json.as_object_mut().and_then(|o| o.remove("memory_key")) {
            let key = key.as_str().unwrap_or("").to_string();
            let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?;
            Ok(Self::Memory { key, message })
        } else {
            let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?;
            Ok(Self::Message(message))
        }
    }
}

impl ConversationEntry {
    /// Get the API message for sending to the model.
    pub fn api_message(&self) -> &Message {
        match self {
            Self::Message(m) => m,
            Self::Memory { message, .. } => message,
        }
    }

    pub fn is_memory(&self) -> bool {
        matches!(self, Self::Memory { .. })
    }

    /// Get a reference to the inner message.
    pub fn message(&self) -> &Message {
        match self {
            Self::Message(m) => m,
            Self::Memory { message, .. } => message,
        }
    }

    /// Get a mutable reference to the inner message.
    pub fn message_mut(&mut self) -> &mut Message {
        match self {
            Self::Message(m) => m,
            Self::Memory { message, .. } => message,
        }
    }
}

#[derive(Clone)]
pub struct ContextState {
    pub system_prompt: String,
    pub personality: Vec<(String, String)>,
    pub journal: Vec<JournalEntry>,
    pub working_stack: Vec<String>,
    /// Conversation entries — messages and memory, interleaved in order.
    /// Does NOT include system prompt, personality, or journal.
    pub entries: Vec<ConversationEntry>,
}

impl ContextState {
    /// Compute the context budget from typed sources.
    pub fn budget(&self, count_str: &dyn Fn(&str) -> usize,
                   count_msg: &dyn Fn(&Message) -> usize,
                   window_tokens: usize) -> ContextBudget {
        let id = count_str(&self.system_prompt)
            + self.personality.iter().map(|(_, c)| count_str(c)).sum::<usize>();
        let jnl: usize = self.journal.iter().map(|e| count_str(&e.content)).sum();
        let mut mem = 0;
        let mut conv = 0;
        for entry in &self.entries {
            let tokens = count_msg(entry.api_message());
            if entry.is_memory() { mem += tokens } else { conv += tokens }
        }
        ContextBudget {
            identity_tokens: id,
            memory_tokens: mem,
            journal_tokens: jnl,
            conversation_tokens: conv,
            window_tokens,
        }
    }

    pub fn render_context_message(&self) -> String {
        let mut parts: Vec<String> = self.personality.iter()
            .map(|(name, content)| format!("## {}\n\n{}", name, content))
            .collect();
        let instructions = std::fs::read_to_string(working_stack::instructions_path()).unwrap_or_default();
        let mut stack_section = instructions;
        if self.working_stack.is_empty() {
            stack_section.push_str("\n## Current stack\n\n(empty)\n");
        } else {
            stack_section.push_str("\n## Current stack\n\n");
            for (i, item) in self.working_stack.iter().enumerate() {
                if i == self.working_stack.len() - 1 {
                    stack_section.push_str(&format!("→ {}\n", item));
                } else {
                    stack_section.push_str(&format!("  [{}] {}\n", i, item));
                }
            }
        }
        parts.push(stack_section);
        parts.join("\n\n---\n\n")
    }
}

#[derive(Debug, Clone, Default)]
pub struct ContextBudget {
    pub identity_tokens: usize,
    pub memory_tokens: usize,
    pub journal_tokens: usize,
    pub conversation_tokens: usize,
    pub window_tokens: usize,
}

impl ContextBudget {
    pub fn used(&self) -> usize {
        self.identity_tokens + self.memory_tokens + self.journal_tokens + self.conversation_tokens
    }
    pub fn free(&self) -> usize {
        self.window_tokens.saturating_sub(self.used())
    }
    pub fn status_string(&self) -> String {
        let total = self.window_tokens;
        if total == 0 { return String::new(); }
        let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / total).max(1) };
        format!("id:{}% mem:{}% jnl:{}% conv:{}% free:{}%",
            pct(self.identity_tokens), pct(self.memory_tokens),
            pct(self.journal_tokens), pct(self.conversation_tokens), pct(self.free()))
    }
}