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 crate::agent::api::*;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use tiktoken_rs::CoreBPE;
use crate::agent::tools::working_stack;

// --- 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, PartialEq)]
pub enum ConversationEntry {
    /// System prompt or system-level instruction.
    System(Message),
    Message(Message),
    Memory { key: String, message: Message, score: Option<f64> },
    /// DMN heartbeat/autonomous prompt — evicted aggressively during compaction.
    Dmn(Message),
    /// Debug/status log line — written to conversation log for tracing,
    /// skipped on read-back.
    Log(String),
}

/// Entry in the context window — wraps a ConversationEntry with cached metadata.
#[derive(Debug, Clone)]
pub struct ContextEntry {
    pub entry: ConversationEntry,
    /// Cached token count (0 for Log entries).
    pub tokens: usize,
    /// When this entry was added to the context.
    pub timestamp: Option<DateTime<Utc>>,
}

/// A named section of the context window with cached token total.
#[derive(Debug, Clone)]
pub struct ContextSection {
    pub name: String,
    /// Cached sum of entry tokens.
    tokens: usize,
    entries: Vec<ContextEntry>,
}

impl ContextSection {
    pub fn new(name: impl Into<String>) -> Self {
        Self { name: name.into(), tokens: 0, entries: Vec::new() }
    }

    pub fn entries(&self) -> &[ContextEntry] { &self.entries }
    pub fn tokens(&self) -> usize { self.tokens }
    pub fn len(&self) -> usize { self.entries.len() }
    pub fn is_empty(&self) -> bool { self.entries.is_empty() }

    /// Push an entry, updating the cached token total.
    pub fn push(&mut self, entry: ContextEntry) {
        self.tokens += entry.tokens;
        self.entries.push(entry);
    }

    /// Replace an entry at `index`, adjusting the token total.
    pub fn set(&mut self, index: usize, entry: ContextEntry) {
        self.tokens -= self.entries[index].tokens;
        self.tokens += entry.tokens;
        self.entries[index] = entry;
    }

    /// Remove an entry at `index`, adjusting the token total.
    pub fn del(&mut self, index: usize) -> ContextEntry {
        let removed = self.entries.remove(index);
        self.tokens -= removed.tokens;
        removed
    }

    /// Replace the message inside an entry, recomputing its token count.
    pub fn set_message(&mut self, index: usize, tokenizer: &CoreBPE, msg: Message) {
        let old_tokens = self.entries[index].tokens;
        *self.entries[index].entry.message_mut() = msg;
        let new_tokens = msg_token_count(tokenizer, self.entries[index].entry.api_message());
        self.entries[index].tokens = new_tokens;
        self.tokens = self.tokens - old_tokens + new_tokens;
    }

    /// Set the score on a Memory entry. No token change.
    pub fn set_score(&mut self, index: usize, score: Option<f64>) {
        if let ConversationEntry::Memory { score: s, .. } = &mut self.entries[index].entry {
            *s = score;
        }
    }

    /// Bulk replace all entries, recomputing token total.
    pub fn set_entries(&mut self, entries: Vec<ContextEntry>) {
        self.tokens = entries.iter().map(|e| e.tokens).sum();
        self.entries = entries;
    }

    /// Dedup and trim entries to fit within context budget.
    pub fn trim(&mut self, budget: &ContextBudget, tokenizer: &CoreBPE) {
        let result = trim_entries(&self.entries, tokenizer, budget);
        self.entries = result;
        self.tokens = self.entries.iter().map(|e| e.tokens).sum();
    }

    /// Clear all entries.
    pub fn clear(&mut self) {
        self.entries.clear();
        self.tokens = 0;
    }
}

#[derive(Clone)]
pub struct ContextState {
    pub system: ContextSection,
    pub identity: ContextSection,
    pub journal: ContextSection,
    pub conversation: ContextSection,
    /// Working stack — separate from identity because it's managed
    /// by its own tool, not loaded from personality files.
    pub working_stack: Vec<String>,
}

impl ContextState {
    /// Total tokens across all sections.
    pub fn total_tokens(&self) -> usize {
        self.system.tokens() + self.identity.tokens()
            + self.journal.tokens() + self.conversation.tokens()
    }

    /// All sections as a slice for iteration.
    pub fn sections(&self) -> [&ContextSection; 4] {
        [&self.system, &self.identity, &self.journal, &self.conversation]
    }
}

/// 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.
fn trim_entries(
    entries: &[ContextEntry],
    _tokenizer: &CoreBPE,
    budget: &ContextBudget,
) -> Vec<ContextEntry> {
    let fixed_tokens = budget.system + budget.identity + budget.journal;
    // --- 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, ce) in entries.iter().enumerate() {
        if let ConversationEntry::Memory { key, .. } = &ce.entry {
            if let Some(prev) = seen_keys.insert(key.as_str(), i) {
                drop_indices.insert(prev);
            }
        }
    }

    let deduped: Vec<ContextEntry> = 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 msg_costs: Vec<usize> = deduped.iter().map(|e| e.tokens).collect();
    let entry_total: usize = msg_costs.iter().sum();
    let total: usize = fixed_tokens + entry_total;

    let mem_tokens: usize = deduped.iter()
        .filter(|ce| ce.entry.is_memory())
        .map(|ce| ce.tokens).sum();
    let conv_tokens: usize = entry_total - mem_tokens;

    dbglog!("[trim] max_tokens={} fixed={} mem={} conv={} total={} entries={}",
        max_tokens, fixed_tokens, mem_tokens, conv_tokens, total, deduped.len());

    // Phase 2a: evict all DMN entries first — they're ephemeral
    let mut drop = vec![false; deduped.len()];
    let mut trimmed = total;
    let mut cur_mem = mem_tokens;

    for i in 0..deduped.len() {
        if deduped[i].entry.is_dmn() {
            drop[i] = true;
            trimmed -= msg_costs[i];
        }
    }

    // Phase 2b: if memories > 50% of context, evict lowest-scored first
    if cur_mem > conv_tokens && trimmed > max_tokens {
        let mut mem_indices: Vec<usize> = (0..deduped.len())
            .filter(|&i| !drop[i] && deduped[i].entry.is_memory())
            .collect();
        mem_indices.sort_by(|&a, &b| {
            let sa = match &deduped[a].entry {
                ConversationEntry::Memory { score, .. } => score.unwrap_or(0.0),
                _ => 0.0,
            };
            let sb = match &deduped[b].entry {
                ConversationEntry::Memory { score, .. } => score.unwrap_or(0.0),
                _ => 0.0,
            };
            sa.partial_cmp(&sb).unwrap_or(std::cmp::Ordering::Equal)
        });
        for i in mem_indices {
            if cur_mem <= conv_tokens { break; }
            if trimmed <= max_tokens { break; }
            drop[i] = true;
            trimmed -= msg_costs[i];
            cur_mem -= msg_costs[i];
        }
    }

    // Phase 2c: if still over, drop oldest conversation entries
    for i in 0..deduped.len() {
        if trimmed <= max_tokens { break; }
        if drop[i] { continue; }
        drop[i] = true;
        trimmed -= msg_costs[i];
    }

    // Walk forward to include complete conversation boundaries
    let mut result: Vec<ContextEntry> = Vec::new();
    let mut skipping = true;
    for (i, ce) in deduped.into_iter().enumerate() {
        if skipping {
            if drop[i] { continue; }
            // Snap to user message boundary
            if ce.entry.message().role != Role::User { continue; }
            skipping = false;
        }
        result.push(ce);
    }

    dbglog!("[trim] result={} trimmed_total={}", result.len(), trimmed);

    result
}

/// 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")
}

// 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::System(m) | Self::Message(m) | Self::Dmn(m) => m.serialize(s),
            Self::Memory { key, message, score } => {
                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)?;
                if let Some(s) = score {
                    map.serialize_entry("memory_score", s)?;
                }
                map.end()
            }
            Self::Log(text) => {
                use serde::ser::SerializeMap;
                let mut map = s.serialize_map(Some(1))?;
                map.serialize_entry("log", text)?;
                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)?;
        // Log entries — skip on read-back
        if json.get("log").is_some() {
            let text = json["log"].as_str().unwrap_or("").to_string();
            return Ok(Self::Log(text));
        }
        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 score = json.as_object_mut()
                .and_then(|o| o.remove("memory_score"))
                .and_then(|v| v.as_f64());
            let message: Message = serde_json::from_value(json).map_err(serde::de::Error::custom)?;
            Ok(Self::Memory { key, message, score })
        } 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.
    /// Panics on Log entries (which should be filtered before API calls).
    pub fn api_message(&self) -> &Message {
        match self {
            Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
            Self::Memory { message, .. } => message,
            Self::Log(_) => panic!("Log entries have no API message"),
        }
    }

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

    pub fn is_dmn(&self) -> bool {
        matches!(self, Self::Dmn(_))
    }

    pub fn is_log(&self) -> bool {
        matches!(self, Self::Log(_))
    }

    /// Get a reference to the inner message.
    /// Panics on Log entries.
    pub fn message(&self) -> &Message {
        match self {
            Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
            Self::Memory { message, .. } => message,
            Self::Log(_) => panic!("Log entries have no message"),
        }
    }

    /// Get a mutable reference to the inner message.
    /// Panics on Log entries.
    pub fn message_mut(&mut self) -> &mut Message {
        match self {
            Self::System(m) | Self::Message(m) | Self::Dmn(m) => m,
            Self::Memory { message, .. } => message,
            Self::Log(_) => panic!("Log entries have no message"),
        }
    }
}

impl ContextState {
    /// Render journal entries into a single text block.
    pub fn render_journal(&self) -> String {
        if self.journal.is_empty() { return String::new(); }
        let mut text = String::from("[Earlier — from your journal]\n\n");
        for e in self.journal.entries() {
            use std::fmt::Write;
            if let Some(ts) = &e.timestamp {
                writeln!(text, "## {}\n{}\n",
                    ts.format("%Y-%m-%dT%H:%M"),
                    e.entry.message().content_text()).ok();
            } else {
                text.push_str(&e.entry.message().content_text());
                text.push_str("\n\n");
            }
        }
        text
    }

    /// Render identity files + working stack into a single user message.
    pub fn render_context_message(&self) -> String {
        let mut parts: Vec<String> = self.identity.entries().iter()
            .map(|e| e.entry.message().content_text().to_string())
            .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")
    }
}

/// Token budget per context category — cheap to compute, no formatting.
pub struct ContextBudget {
    pub system: usize,
    pub identity: usize,
    pub journal: usize,
    pub memory: usize,
    pub conversation: usize,
}

impl ContextBudget {
    pub fn total(&self) -> usize {
        self.system + self.identity + self.journal + self.memory + self.conversation
    }

    pub fn format(&self) -> String {
        let window = context_window();
        if window == 0 { return String::new(); }
        let used = self.total();
        let free = window.saturating_sub(used);
        let pct = |n: usize| if n == 0 { 0 } else { ((n * 100) / window).max(1) };
        format!("sys:{}% id:{}% jnl:{}% mem:{}% conv:{}% free:{}%",
            pct(self.system), pct(self.identity), pct(self.journal),
            pct(self.memory), pct(self.conversation), pct(free))
    }
}