Conversation Pruner — Rust source

//! Conversation Pruner — compute a deterministic pruning plan for a
//! token-budgeted chat history.
//!
//! Language: Rust (edition 2021, standard library only)
//! Source:   CosmoDev polyglot showcase port of the Conversation Pruner
//!           tool (slug: conversation-pruner).
//! Port of src/lib/conversationPruner.ts (the canonical TypeScript
//!           implementation).
//! Tool page: https://dev.cosmolabs.org/tools/conversation-pruner
//! License:  display source — part of CosmoDev's polyglot tool pages.
//!
//! Given per-message token counts and a context budget, decide which
//! messages to keep verbatim, which to fold into one running summary, and
//! which to drop outright — protecting system messages, pinned turns, the
//! first turn (the opening user request), and the current (last user)
//! request. Token counts are plain integers; the only float math is the
//! summary cost: fixed framing plus 10% of the folded content, rounded up.

use std::fmt;

/// Summary compression model: fixed framing tokens.
pub const SUMMARY_FIXED_TOKENS: i64 = 60;

/// Summary compression model: share of the folded content's tokens.
pub const SUMMARY_RATIO: f64 = 0.1;

/// A chat participant. Mirrors the TypeScript `ChatRole` union.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Role {
    System,
    User,
    Assistant,
    Tool,
}

impl Role {
    /// The lowercase wire name used by the TypeScript union.
    pub fn as_str(&self) -> &'static str {
        match self {
            Role::System => "system",
            Role::User => "user",
            Role::Assistant => "assistant",
            Role::Tool => "tool",
        }
    }
}

/// One message plus its prompt-side token count. Pinned messages are never
/// dropped or summarized (the TypeScript's optional flag defaults to false).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConversationMessage {
    pub role: Role,
    pub content: String,
    /// Token count for this message (prompt-side framing included by caller).
    pub tokens: i64,
    pub pinned: bool,
}

/// What the plan does with one message.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PruneAction {
    Keep,
    Summarize,
    Drop,
}

impl PruneAction {
    /// The lowercase action name used by the TypeScript union.
    pub fn as_str(&self) -> &'static str {
        match self {
            PruneAction::Keep => "keep",
            PruneAction::Summarize => "summarize",
            PruneAction::Drop => "drop",
        }
    }
}

/// The plan's verdict on one message.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PruneDecision {
    pub index: usize,
    pub role: Role,
    pub action: PruneAction,
    pub tokens: i64,
}

/// The full pruning plan for one conversation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PrunePlan {
    pub decisions: Vec<PruneDecision>,
    pub kept_tokens: i64,
    pub summarized_tokens: i64,
    pub dropped_tokens: i64,
    /// Tokens the summary placeholder itself will cost in the prompt.
    pub summary_cost_tokens: i64,
    pub projected_tokens: i64,
    pub fits_budget: bool,
    pub warnings: Vec<String>,
}

/// Validation failures. The TypeScript throws a `RangeError`; Rust returns one.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PruneError {
    /// `budget_tokens` was negative.
    NegativeBudget,
    /// Some message carried a negative token count.
    NegativeMessageTokens,
}

impl fmt::Display for PruneError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PruneError::NegativeBudget => write!(f, "budgetTokens must be >= 0"),
            PruneError::NegativeMessageTokens => write!(f, "message tokens must be >= 0"),
        }
    }
}

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

/// Group an integer with en-US commas: `1234567` → `"1,234,567"`, matching
/// `toLocaleString('en-US')` in the TypeScript.
fn thousands(n: i64) -> String {
    let digits = n.unsigned_abs().to_string();
    let len = digits.len();
    let mut out = String::with_capacity(len + len / 3 + 1);
    if n < 0 {
        out.push('-');
    }
    for (i, ch) in digits.chars().enumerate() {
        if i > 0 && (len - i) % 3 == 0 {
            out.push(',');
        }
        out.push(ch);
    }
    out
}

/// Compute the pruning plan for `messages` under `budget_tokens`.
pub fn plan_prune(
    messages: &[ConversationMessage],
    budget_tokens: i64,
) -> Result<PrunePlan, PruneError> {
    let mut warnings: Vec<String> = Vec::new();
    if budget_tokens < 0 {
        return Err(PruneError::NegativeBudget);
    }
    if messages.iter().any(|m| m.tokens < 0) {
        return Err(PruneError::NegativeMessageTokens);
    }

    let n = messages.len();
    let last_user: Option<usize> = (0..n).rev().find(|&i| messages[i].role == Role::User);

    // Untouchable: every system message, pinned messages, the first turn (the
    // opening user request that anchors the conversation), and the current
    // request (the last user message and everything after it).
    let mut protected = vec![false; n];
    for (i, m) in messages.iter().enumerate() {
        if m.role == Role::System || m.pinned {
            protected[i] = true;
        }
    }
    if n > 0 {
        protected[0] = true;
    }
    if let Some(first_turn) = (0..n).find(|&i| messages[i].role != Role::System) {
        protected[first_turn] = true;
    }
    let tail_start = match last_user {
        Some(i) => i,
        // No user turn: only the final message is "current".
        None => n.saturating_sub(1),
    };
    for p in protected.iter_mut().skip(tail_start) {
        *p = true;
    }

    let protected_tokens: i64 = (0..n)
        .filter(|&i| protected[i])
        .map(|i| messages[i].tokens)
        .sum();
    if protected_tokens > budget_tokens {
        warnings.push(format!(
            "Protected messages alone are {} tokens against a {} budget — raise the budget (or reserve less for the reply) before pruning anything else.",
            thousands(protected_tokens),
            thousands(budget_tokens)
        ));
    }

    // Fill the remaining budget newest-to-oldest through the middle.
    let mut actions = vec![PruneAction::Drop; n];
    for i in 0..n {
        if protected[i] {
            actions[i] = PruneAction::Keep;
        }
    }
    let mut used = protected_tokens;
    for i in (0..n).rev() {
        if actions[i] != PruneAction::Drop {
            continue;
        }
        if used + messages[i].tokens <= budget_tokens {
            actions[i] = PruneAction::Keep;
            used += messages[i].tokens;
        } else {
            break; // oldest-unfilled remain drop/summarize candidates, newest first stopped
        }
    }

    // Everything still 'drop' in the middle folds into ONE running summary when
    // the compressed form fits where the raw turns did not.
    let summarize_idx: Vec<usize> = (0..n)
        .filter(|&i| actions[i] == PruneAction::Drop && !protected[i])
        .collect();
    let summarize_tokens: i64 = summarize_idx.iter().map(|&i| messages[i].tokens).sum();
    let attempted_summary_cost = if summarize_idx.is_empty() {
        0
    } else {
        SUMMARY_FIXED_TOKENS + (summarize_tokens as f64 * SUMMARY_RATIO).ceil() as i64
    };

    // The summary only costs anything when it is actually applied — otherwise
    // those turns drop and cost zero.
    let mut summary_cost = 0;
    if attempted_summary_cost > 0 && used + attempted_summary_cost <= budget_tokens {
        for &i in &summarize_idx {
            actions[i] = PruneAction::Summarize;
        }
        summary_cost = attempted_summary_cost;
        // Mirrors the TypeScript; `used` is not consulted again afterwards.
        #[allow(unused_assignments)]
        {
            used += summary_cost;
        }
    } else if attempted_summary_cost > 0 {
        warnings.push(format!(
            "Even the compressed summary ({} tokens) does not fit the remaining budget — the oldest turns are dropped instead.",
            thousands(attempted_summary_cost)
        ));
    }

    let decisions: Vec<PruneDecision> = (0..n)
        .map(|i| PruneDecision {
            index: i,
            role: messages[i].role,
            action: actions[i],
            tokens: messages[i].tokens,
        })
        .collect();

    let mut kept_tokens = 0;
    let mut dropped_tokens = 0;
    let mut folded_tokens = 0;
    for d in &decisions {
        match d.action {
            PruneAction::Keep => kept_tokens += d.tokens,
            PruneAction::Drop => dropped_tokens += d.tokens,
            PruneAction::Summarize => folded_tokens += d.tokens,
        }
    }

    Ok(PrunePlan {
        decisions,
        kept_tokens,
        summarized_tokens: folded_tokens,
        dropped_tokens,
        summary_cost_tokens: summary_cost,
        projected_tokens: kept_tokens + summary_cost,
        fits_budget: kept_tokens + summary_cost <= budget_tokens,
        warnings,
    })
}

/// Human-readable one-line summary of a plan.
pub fn describe_prune(plan: &PrunePlan) -> String {
    if !plan.fits_budget {
        return format!(
            "Does not fit: {} tokens projected against the budget.",
            thousands(plan.projected_tokens)
        );
    }
    let mut parts = vec![format!("{} kept", thousands(plan.kept_tokens))];
    if plan.summarized_tokens > 0 {
        parts.push(format!(
            "{} folded into a {}-token summary",
            thousands(plan.summarized_tokens),
            thousands(plan.summary_cost_tokens)
        ));
    }
    if plan.dropped_tokens > 0 {
        parts.push(format!("{} dropped", thousands(plan.dropped_tokens)));
    }
    format!("{} — fits the budget.", parts.join(" · "))
}