DS4Server/src/compaction.rs

use crate::engine::ChatTurn;

pub(crate) const SUMMARY_MAX_TOKENS: i32 = 4096;
pub(crate) const TOOL_RESULT_RESERVE_TOKENS: u32 = 1024;
const MIN_SUMMARY_TOKENS: u32 = 256;

const SOFT_PERCENT: u32 = 85;
const MIN_FREE_TOKENS: u32 = 8192;
const TAIL_DIVISOR: u32 = 10;
const TAIL_CAP_TOKENS: u32 = 50_000;

pub(crate) const SUMMARY_PREFIX: &str =
    "[DS4Server compacted earlier conversation. Durable task-state summary follows.]";
pub(crate) const SUMMARY_SUFFIX: &str =
    "[End compacted summary. Recent conversation continues verbatim below.]";

pub(crate) fn should_compact(used: u32, limit: u32) -> bool {
    if used == 0 || limit == 0 {
        return false;
    }
    used >= limit.saturating_mul(SOFT_PERCENT) / 100
        || limit.saturating_sub(used) <= MIN_FREE_TOKENS.min(limit / 8)
}

pub(crate) fn tool_result_reserve(limit: u32) -> u32 {
    TOOL_RESULT_RESERVE_TOKENS.min((limit / 8).max(16))
}

pub(crate) fn tool_result_fits(projected: u32, limit: u32, reserve: u32) -> bool {
    limit > 0 && projected.saturating_add(reserve) < limit
}

pub(crate) fn bounded_tool_error(projected: u32, limit: u32, reserve: u32) -> String {
    format!(
        "Tool error: tool result still does not fit after context compaction (projected_prompt={projected} tokens, ctx={limit}, reserve={reserve}). Retry with a smaller read/search/bash output.\n"
    )
}

pub(crate) fn summary_budget(prompt: u32, limit: u32) -> Option<i32> {
    let room = limit.saturating_sub(prompt).saturating_sub(1);
    (room >= MIN_SUMMARY_TOKENS).then_some(room.min(SUMMARY_MAX_TOKENS as u32) as i32)
}

pub(crate) fn tail_budget(context: u32) -> u32 {
    (context / TAIL_DIVISOR).clamp(1, TAIL_CAP_TOKENS)
}

/// Matches the reference token-level rule: start at the target unless a user
/// boundary appears between it and the end. Message starts are token offsets.
pub(crate) fn tail_start(messages: &[ChatTurn], starts: &[u32], bottom: u32, budget: u32) -> usize {
    let target = bottom.saturating_sub(budget);
    messages
        .iter()
        .zip(starts)
        .position(|(message, start)| message.user && *start >= target)
        .or_else(|| starts.iter().position(|start| *start >= target))
        .unwrap_or(messages.len())
}

pub(crate) fn summary_prompt(reason: &str) -> String {
    format!(
        "Internal DS4Server context compaction request. This is not a user request.\n\
Write a durable task-state summary of the conversation so far. Preserve only facts that matter for continuing the work:\n\
- user goals, constraints, and preferences\n\
- files inspected or edited\n\
- commands run and important results\n\
- decisions, rejected approaches, known bugs, and pending next steps\n\
- reloadable bulky data with exact paths/ranges/commands when available\n\n\
Do not invent facts. Do not include generic narration or raw file contents unless essential.\n\
After the summary, stop. Do not continue the user task, call tools, or output thinking/DSML control markup.\n\
Output only the compact summary.\n\nCompaction reason: {reason}\n"
    )
}

pub(crate) fn summary_system_prompt(base: &str, summary: Option<&str>) -> String {
    match summary.filter(|summary| !summary.trim().is_empty()) {
        Some(summary) => format!(
            "{base}\n\n{SUMMARY_PREFIX}\n{}\n{SUMMARY_SUFFIX}",
            summary.trim()
        ),
        None => base.to_owned(),
    }
}

pub(crate) fn sanitize_summary(summary: &str) -> String {
    let end = ["<｜DSML｜", "<DSML｜", "<tool_call>", "<think>", "</think>"]
        .into_iter()
        .filter_map(|marker| summary.find(marker))
        .min()
        .unwrap_or(summary.len());
    summary[..end].trim().to_owned()
}

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

    fn turn(user: bool) -> ChatTurn {
        ChatTurn {
            user,
            tool: false,
            system: false,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: String::new(),
        }
    }

    #[test]
    fn soft_trigger_matches_reference_thresholds() {
        assert!(should_compact(85_000, 100_000));
        assert!(should_compact(92_000, 100_000));
        assert!(!should_compact(84_999, 100_000));
        assert!(should_compact(6_963, 8_192));
        assert!(!should_compact(6_962, 8_192));
    }

    #[test]
    fn tail_prefers_the_first_user_boundary_after_target() {
        let messages = [turn(true), turn(false), turn(true), turn(false)];
        assert_eq!(tail_start(&messages, &[10, 30, 70, 90], 100, 40), 2);
        assert_eq!(tail_start(&messages, &[10, 30, 70, 90], 100, 15), 3);
    }

    #[test]
    fn summary_control_markup_is_private() {
        assert_eq!(
            sanitize_summary("state kept\n<tool_call>bash"),
            "state kept"
        );
    }

    #[test]
    fn a_new_summary_replaces_the_previous_rebuild_summary() {
        let current = summary_system_prompt("tools", Some("old state"));
        assert!(current.contains("old state"));
        let rebuilt = summary_system_prompt("tools", Some("new state"));
        assert!(rebuilt.contains("new state"));
        assert!(!rebuilt.contains("old state"));
    }

    #[test]
    fn hard_trigger_and_retry_reserve_match_the_reference() {
        assert_eq!(tool_result_reserve(4096), 512);
        assert_eq!(tool_result_reserve(65_536), 1024);
        assert!(tool_result_fits(3000, 4096, 512));
        assert!(!tool_result_fits(3584, 4096, 512));
        let error = bounded_tool_error(5000, 4096, 512);
        assert!(error.starts_with("Tool error:"));
        assert!(error.contains("projected_prompt=5000 tokens"));
        assert!(error.len() < 256);
    }

    #[test]
    fn summary_budget_keeps_answer_room_and_rejects_exhausted_contexts() {
        assert_eq!(summary_budget(1000, 8192), Some(SUMMARY_MAX_TOKENS));
        assert_eq!(summary_budget(7900, 8192), Some(291));
        assert_eq!(summary_budget(7936, 8192), None);
    }
}