Add agent context compaction

src/compaction.rs

@@ -0,0 +1,141 @@
+use crate::engine::ChatTurn;
+
+pub(crate) const SUMMARY_MAX_TOKENS: i32 = 4096;
+pub(crate) const TOOL_RESULT_RESERVE_TOKENS: u32 = 1024;
+
+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_needs_compaction(used: u32, limit: u32, result: &str) -> bool {
+    should_compact(used, limit)
+        || used
+            .saturating_add(result.len().min(u32::MAX as usize) as u32)
+            .saturating_add(TOOL_RESULT_RESERVE_TOKENS)
+            >= limit
+}
+
+pub(crate) fn bounded_tool_result(used: u32, limit: u32, result: String) -> String {
+    if used
+        .saturating_add(result.len().min(u32::MAX as usize) as u32)
+        .saturating_add(TOOL_RESULT_RESERVE_TOKENS)
+        < limit
+    {
+        result
+    } else {
+        "Tool error: the result is too large for the remaining context after compaction. Retry with a smaller read/search/bash output.\n".into()
+    }
+}
+
+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,
+            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 oversized_tool_result_becomes_a_bounded_retry_error() {
+        let result = "x".repeat(4_000);
+        assert!(tool_result_needs_compaction(6_000, 10_000, &result));
+        let error = bounded_tool_result(4_000, 5_000, result);
+        assert!(error.starts_with("Tool error:"));
+        assert!(error.len() < 160);
+    }
+}