boxmaker/crates/loopd/src/compact.rs

//! Compaction: when a session starts a new epoch, the summary call, and writing the next files
//! (M5d spec, section 3). Task 07 writes the next files; task 08 calls this after a turn.

use bxcore::sanitize::neutralise;
use bxcore::{Epoch, LogRecord, Timestamp, Usage};

use crate::baseline::{Baseline, messages_with_carry};
use crate::config::Config;
use crate::llama::{ChatMessage, ChatRequest, Client};
use crate::memory::MemorySource;
use crate::session::Session;
use crate::tools::Registry;
use crate::turn::is_context_full;
use crate::turn::{Runtime, TurnError, TurnOutcome, run_turn};
use bxcore::TurnEvent;

/// Why a summary call or a next-epoch write failed. Every variant's `Display` ends with the
/// runbook pointer `session-compact-failed`.
#[derive(Debug)]
pub enum CompactError {
    Failed { detail: String },
}

impl std::fmt::Display for CompactError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CompactError::Failed { detail } => {
                write!(f, "{detail}; see docs/runbook.md#session-compact-failed")
            }
        }
    }
}

impl std::error::Error for CompactError {}

/// True when `cache_n + prompt_n` is at or above `threshold * n_ctx`.
pub fn due(usage: &Usage, threshold: f64, n_ctx: u64) -> bool {
    let total = usage.cache_n.saturating_add(usage.prompt_n) as f64;
    total >= threshold * (n_ctx as f64)
}

/// The epoch's `User` texts, newest last, joined with `"\n\n"`, then `cap_chars`.
pub fn fallback_summary(records: &[LogRecord], summary_chars: usize) -> String {
    let contents: Vec<&str> = records
        .iter()
        .filter_map(|r| match r {
            LogRecord::User { content, .. } => Some(content.as_str()),
            _ => None,
        })
        .collect();
    let joined = contents
        .iter()
        .rev()
        .cloned()
        .collect::<Vec<_>>()
        .join("\n\n");
    cap_chars(&joined, summary_chars)
}

pub(crate) fn cap_chars(s: &str, n: usize) -> String {
    s.chars().take(n).collect()
}

/// One completion on the background slot: no tools, thinking on, last user `compact.md`.
/// Context-full uses `fallback_summary`. A missing prompt, or any other infer error, is
/// `CompactError::Failed` with the runbook pointer. No fallback for a missing prompt.
pub fn summarise(cfg: &Config, client: &Client, session: &Session) -> Result<String, CompactError> {
    let path = cfg.compact_prompt();
    let prompt = std::fs::read_to_string(&path).map_err(|e| CompactError::Failed {
        detail: format!("{}: {e}", path.display()),
    })?;
    let mut messages = messages_with_carry(session.baseline(), session.records(), session.carry());
    messages.push(ChatMessage::User {
        content: prompt.trim_end().to_string(),
    });
    let request = ChatRequest {
        slot: cfg.slots.background,
        messages,
        tools: Vec::new(),
        thinking: true,
    };
    match client.chat_with_retry(&request, &mut |_| {}) {
        Ok(completion) => {
            let text = completion.content.unwrap_or_default();
            let text = neutralise(&text);
            Ok(cap_chars(text.as_ref(), cfg.compact.summary_chars))
        }
        Err(e) if is_context_full(&e) => Ok(fallback_summary(
            session.records(),
            cfg.compact.summary_chars,
        )),
        Err(e) => Err(CompactError::Failed {
            detail: e.to_string(),
        }),
    }
}

/// Append `EpochEnd` (unless it is already last), rebuild the baseline from `memory`, and
/// switch `session` onto the next epoch. When the current log already ends with `EpochEnd`,
/// that record's summary is the carry and the argument is ignored.
pub fn commit(
    session: &mut Session,
    cfg: &Config,
    registry: &Registry,
    memory: &dyn MemorySource,
    summary: String,
) -> Result<(), CompactError> {
    let (next, summary) = match session.records().last() {
        Some(LogRecord::EpochEnd { next, summary, .. }) => (*next, summary.clone()),
        _ => {
            let next = Epoch(session.epoch().0.saturating_add(1));
            session
                .append(LogRecord::EpochEnd {
                    time: Timestamp::now(),
                    next,
                    summary: summary.clone(),
                })
                .map_err(|e| CompactError::Failed {
                    detail: e.to_string(),
                })?;
            (next, summary)
        }
    };
    let block = memory
        .block()
        .map_err(|e| CompactError::Failed { detail: e })?;
    let baseline =
        Baseline::assemble_with(cfg, registry, &block).map_err(|e| CompactError::Failed {
            detail: e.to_string(),
        })?;
    session
        .begin_epoch(next, baseline, summary)
        .map_err(|e| CompactError::Failed {
            detail: e.to_string(),
        })?;
    Ok(())
}

/// After `run_turn`: compact when the last `Usage` is due, or on `SessionFull` (then retry the
/// user message once). Compact failure on a successful turn keeps the success; on `SessionFull`
/// keeps `SessionFull`. Other errors skip compact.
pub fn finish_turn(
    session: &mut Session,
    rt: &Runtime<'_>,
    memory: &dyn MemorySource,
    text: &str,
    speaker: Option<&str>,
    on_event: &mut dyn FnMut(&TurnEvent),
    outcome: Result<TurnOutcome, TurnError>,
) -> Result<TurnOutcome, TurnError> {
    match outcome {
        Err(TurnError::SessionFull) => {
            if let Err(e) = crate::summarise::missing(rt.cfg, rt.client, session) {
                eprintln!("loopd: tool summary: {e}");
            }
            match compact_session(session, rt, memory) {
                Ok(()) => run_turn(session, rt, text, speaker, on_event),
                Err(e) => {
                    log_compact(session, &e);
                    Err(TurnError::SessionFull)
                }
            }
        }
        Ok(done) => {
            if let Err(e) = crate::summarise::missing(rt.cfg, rt.client, session) {
                eprintln!("loopd: tool summary: {e}");
            }
            if let Some(usage) = session.last_usage()
                && due(&usage, rt.cfg.compact.threshold, rt.cfg.expect.n_ctx)
                && let Err(e) = compact_session(session, rt, memory)
            {
                log_compact(session, &e);
            }
            Ok(done)
        }
        Err(e) => Err(e),
    }
}

fn compact_session(
    session: &mut Session,
    rt: &Runtime<'_>,
    memory: &dyn MemorySource,
) -> Result<(), CompactError> {
    let summary = summarise(rt.cfg, rt.client, session)?;
    commit(session, rt.cfg, rt.registry, memory, summary)
}

fn log_compact(session: &Session, e: &CompactError) {
    eprintln!(
        "loopd: could not compact session {}: {e}",
        session.id().as_str()
    );
}